Skip to content

Commit

Permalink
feat: Add option to find, load and merge stubs-only packages
Browse files Browse the repository at this point in the history
Add an option to look for X-stubs first when looking for the X module.
Stubs-only packages are described in PEP 561.

PR #221: #221
Co-authored-by: Timothée Mazzucotelli <pawamoy@pm.me>
  • Loading branch information
romain-intel authored Nov 10, 2023
1 parent b7048d0 commit 6e55f3b
Show file tree
Hide file tree
Showing 7 changed files with 217 additions and 8 deletions.
19 changes: 18 additions & 1 deletion src/griffe/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@ def _load_packages(
resolve_external: bool = False,
allow_inspection: bool = True,
store_source: bool = True,
find_stubs_package: bool = False,
) -> GriffeLoader:
# Create a single loader.
loader = GriffeLoader(
Expand All @@ -92,7 +93,7 @@ def _load_packages(
continue
logger.info(f"Loading package {package}")
try:
loader.load_module(package)
loader.load_module(package, try_relative_path=True, find_stubs_package=find_stubs_package)
except ModuleNotFoundError as error:
logger.error(f"Could not find package {package}: {error}") # noqa: TRY400
except ImportError as error:
Expand Down Expand Up @@ -153,6 +154,14 @@ def add_common_options(subparser: argparse.ArgumentParser) -> None:
help="Paths to search packages into.",
)
loading_options = subparser.add_argument_group(title="Loading options")
loading_options.add_argument(
"-B",
"--find-stubs-packages",
dest="find_stubs_package",
action="store_true",
default=False,
help="Whether to look for stubs-only packages and merge them with concrete ones.",
)
loading_options.add_argument(
"-e",
"--extensions",
Expand Down Expand Up @@ -306,6 +315,7 @@ def dump(
resolve_implicit: bool = False,
resolve_external: bool = False,
search_paths: Sequence[str | Path] | None = None,
find_stubs_package: bool = False,
append_sys_path: bool = False,
allow_inspection: bool = True,
stats: bool = False,
Expand All @@ -323,6 +333,9 @@ def dump(
resolve_external: Whether to load additional, unspecified modules to resolve aliases.
extensions: The extensions to use.
search_paths: The paths to search into.
find_stubs_package: Whether to search for stubs-only packages.
If both the package and its stubs are found, they'll be merged together.
If only the stubs are found, they'll be used as the package itself.
append_sys_path: Whether to append the contents of `sys.path` to the search paths.
allow_inspection: Whether to allow inspecting modules when visiting them is not possible.
stats: Whether to compute and log stats about loading.
Expand Down Expand Up @@ -357,6 +370,7 @@ def dump(
resolve_external=resolve_external,
allow_inspection=allow_inspection,
store_source=False,
find_stubs_package=find_stubs_package,
)
data_packages = loader.modules_collection.members

Expand Down Expand Up @@ -385,6 +399,7 @@ def check(
base_ref: str | None = None,
extensions: Sequence[str | dict[str, Any] | ExtensionType | type[ExtensionType]] | None = None,
search_paths: Sequence[str | Path] | None = None,
find_stubs_package: bool = False,
allow_inspection: bool = True,
verbose: bool = False,
color: bool | None = None,
Expand Down Expand Up @@ -438,6 +453,7 @@ def check(
extensions=loaded_extensions,
search_paths=search_paths,
allow_inspection=allow_inspection,
find_stubs_package=find_stubs_package,
)
else:
new_package = load(
Expand All @@ -446,6 +462,7 @@ def check(
extensions=loaded_extensions,
search_paths=search_paths,
allow_inspection=allow_inspection,
find_stubs_package=find_stubs_package,
)

# Find and display API breakages.
Expand Down
45 changes: 42 additions & 3 deletions src/griffe/finder.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ def find_spec(
module: str | Path,
*,
try_relative_path: bool = True,
find_stubs_package: bool = False,
) -> tuple[str, Package | NamespacePackage]:
"""Find the top module of a module.
Expand All @@ -113,6 +114,9 @@ def find_spec(
module: The module name or path.
try_relative_path: Whether to try finding the module as a relative path,
when the given module is not already a path.
find_stubs_package: Whether to search for stubs-only package.
If both the package and its stubs are found, they'll be merged together.
If only the stubs are found, they'll be used as the package itself.
Raises:
FileNotFoundError: When a Path was passed and the module could not be found:
Expand Down Expand Up @@ -146,7 +150,35 @@ def find_spec(
else:
module_name = module
top_module_name = module.split(".", 1)[0]
return module_name, self.find_package(top_module_name)

# Only search for actual package, let exceptions bubble up.
if not find_stubs_package:
return module_name, self.find_package(top_module_name)

# Search for both package and stubs-only package.
try:
package = self.find_package(top_module_name)
except ModuleNotFoundError:
package = None
try:
stubs = self.find_package(top_module_name + "-stubs")
except ModuleNotFoundError:
stubs = None

# None found, raise error.
if package is None and stubs is None:
raise ModuleNotFoundError(top_module_name)

# Both found, assemble them to be merged later.
if package and stubs:
if isinstance(package, Package) and isinstance(stubs, Package):
package.stubs = stubs.path
elif isinstance(package, NamespacePackage) and isinstance(stubs, NamespacePackage):
package.path += stubs.path
return module_name, package

# Return either one.
return module_name, package or stubs # type: ignore[return-value]

def find_package(self, module_name: str) -> Package | NamespacePackage:
"""Find a package or namespace package.
Expand All @@ -168,6 +200,9 @@ def find_package(self, module_name: str) -> Package | NamespacePackage:
Path(f"{module_name}.py"),
]

real_module_name = module_name
if real_module_name.endswith("-stubs"):
real_module_name = real_module_name[:-6]
namespace_dirs = []
for path in self.search_paths:
path_contents = self._contents(path)
Expand All @@ -177,11 +212,15 @@ def find_package(self, module_name: str) -> Package | NamespacePackage:
if abs_path in path_contents:
if abs_path.suffix:
stubs = abs_path.with_suffix(".pyi")
return Package(module_name, abs_path, stubs if stubs.exists() else None)
return Package(real_module_name, abs_path, stubs if stubs.exists() else None)
init_module = abs_path / "__init__.py"
if init_module.exists() and not _is_pkg_style_namespace(init_module):
stubs = init_module.with_suffix(".pyi")
return Package(module_name, init_module, stubs if stubs.exists() else None)
return Package(real_module_name, init_module, stubs if stubs.exists() else None)
init_module = abs_path / "__init__.pyi"
if init_module.exists():
# Stubs package
return Package(real_module_name, init_module, None)
namespace_dirs.append(abs_path)

if namespace_dirs:
Expand Down
5 changes: 5 additions & 0 deletions src/griffe/git.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,7 @@ def load_git(
lines_collection: LinesCollection | None = None,
modules_collection: ModulesCollection | None = None,
allow_inspection: bool = True,
find_stubs_package: bool = False,
) -> Module:
"""Load and return a module from a specific Git reference.
Expand All @@ -147,6 +148,9 @@ def load_git(
lines_collection: A collection of source code lines.
modules_collection: A collection of modules.
allow_inspection: Whether to allow inspecting modules when visiting them is not possible.
find_stubs_package: Whether to search for stubs-only package.
If both the package and its stubs are found, they'll be merged together.
If only the stubs are found, they'll be used as the package itself.
Returns:
A loaded module.
Expand All @@ -166,6 +170,7 @@ def load_git(
lines_collection=lines_collection,
modules_collection=modules_collection,
allow_inspection=allow_inspection,
find_stubs_package=find_stubs_package,
)


Expand Down
22 changes: 20 additions & 2 deletions src/griffe/loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,13 +93,17 @@ def load_module(
*,
submodules: bool = True,
try_relative_path: bool = True,
find_stubs_package: bool = False,
) -> Module:
"""Load a module.
Parameters:
module: The module name or path.
submodules: Whether to recurse on the submodules.
try_relative_path: Whether to try finding the module as a relative path.
find_stubs_package: Whether to search for stubs-only package.
If both the package and its stubs are found, they'll be merged together.
If only the stubs are found, they'll be used as the package itself.
Raises:
LoadingError: When loading a module failed for various reasons.
Expand All @@ -119,7 +123,11 @@ def load_module(
return self.modules_collection.get_member(module_name)
raise LoadingError("Cannot load builtin module without inspection")
try:
module_name, package = self.finder.find_spec(module, try_relative_path=try_relative_path)
module_name, package = self.finder.find_spec(
module,
try_relative_path=try_relative_path,
find_stubs_package=find_stubs_package,
)
except ModuleNotFoundError:
logger.debug(f"Could not find {module}")
if self.allow_inspection:
Expand Down Expand Up @@ -429,7 +437,12 @@ def _load_package(self, package: Package | NamespacePackage, *, submodules: bool
return top_module
if package.stubs:
self.expand_wildcards(top_module)
stubs = self._load_module(package.name, package.stubs, submodules=False)
# If stubs are in the package itself, they have been merged while loading modules,
# so only the top-level init module needs to be merged still.
# If stubs are in another package (a stubs-only package),
# then we need to load the entire stubs package to merge everything.
submodules = submodules and package.stubs.parent != package.path.parent
stubs = self._load_module(package.name, package.stubs, submodules=submodules)
return merge_stubs(top_module, stubs)
return top_module

Expand Down Expand Up @@ -622,6 +635,7 @@ def load(
lines_collection: LinesCollection | None = None,
modules_collection: ModulesCollection | None = None,
allow_inspection: bool = True,
find_stubs_package: bool = False,
) -> Module:
"""Load and return a module.
Expand Down Expand Up @@ -654,6 +668,9 @@ def load(
lines_collection: A collection of source code lines.
modules_collection: A collection of modules.
allow_inspection: Whether to allow inspecting modules when visiting them is not possible.
find_stubs_package: Whether to search for stubs-only package.
If both the package and its stubs are found, they'll be merged together.
If only the stubs are found, they'll be used as the package itself.
Returns:
A loaded module.
Expand All @@ -670,6 +687,7 @@ def load(
module=module,
submodules=submodules,
try_relative_path=try_relative_path,
find_stubs_package=find_stubs_package,
)


Expand Down
7 changes: 7 additions & 0 deletions src/griffe/merger.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,12 +56,19 @@ def _merge_stubs_overloads(obj: Module | Class, stubs: Module | Class) -> None:
def _merge_stubs_members(obj: Module | Class, stubs: Module | Class) -> None:
for member_name, stub_member in stubs.members.items():
if member_name in obj.members:
# We don't merge imported stub objects that already exist in the concrete module.
# Stub objects must be defined where they are exposed in the concrete package,
# not be imported from other stub modules.
if stub_member.is_alias:
continue
obj_member = obj.get_member(member_name)
with suppress(AliasResolutionError, CyclicAliasError):
if obj_member.kind is not stub_member.kind:
logger.debug(
f"Cannot merge stubs for {obj_member.path}: kind {stub_member.kind.value} != {obj_member.kind.value}",
)
elif obj_member.is_module:
_merge_module_stubs(obj_member, stub_member) # type: ignore[arg-type]
elif obj_member.is_class:
_merge_class_stubs(obj_member, stub_member) # type: ignore[arg-type]
elif obj_member.is_function:
Expand Down
85 changes: 84 additions & 1 deletion tests/test_finder.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

import pytest

from griffe.finder import ModuleFinder, NamespacePackage, _handle_editable_module, _handle_pth_file
from griffe.finder import ModuleFinder, NamespacePackage, Package, _handle_editable_module, _handle_pth_file
from griffe.tests import temporary_pypackage


Expand Down Expand Up @@ -164,3 +164,86 @@ def test_scikit_build_core_file_handling(tmp_path: Path) -> None:
# (they don't respect standard package or namespace package layouts,
# and rely on dynamic meta path finder stuff)
assert _handle_editable_module(pth_file) == [Path("/path/to/whatever")]


@pytest.mark.parametrize(
("first", "second", "find_stubs", "expect"),
[
("package", "stubs", True, "both"),
("stubs", "package", True, "both"),
("package", None, True, "package"),
(None, "package", True, "package"),
("stubs", None, True, "stubs"),
(None, "stubs", True, "stubs"),
(None, None, True, "none"),
("package", "stubs", False, "package"),
("stubs", "package", False, "package"),
("package", None, False, "package"),
(None, "package", False, "package"),
("stubs", None, False, "none"),
(None, "stubs", False, "none"),
(None, None, False, "none"),
],
)
def test_finding_stubs_packages(
tmp_path: Path,
first: str | None,
second: str | None,
find_stubs: bool,
expect: str,
) -> None:
"""Find stubs-only packages.
Parameters:
tmp_path: Pytest fixture.
"""
search_path1 = tmp_path / "sp1"
search_path2 = tmp_path / "sp2"
search_path1.mkdir()
search_path2.mkdir()

if first == "package":
package = search_path1 / "package"
package.mkdir()
package.joinpath("__init__.py").touch()
elif first == "stubs":
stubs = search_path1 / "package-stubs"
stubs.mkdir()
stubs.joinpath("__init__.pyi").touch()

if second == "package":
package = search_path2 / "package"
package.mkdir()
package.joinpath("__init__.py").touch()
elif second == "stubs":
stubs = search_path2 / "package-stubs"
stubs.mkdir()
stubs.joinpath("__init__.pyi").touch()

finder = ModuleFinder([search_path1, search_path2])

if expect == "none":
with pytest.raises(ModuleNotFoundError):
finder.find_spec("package", try_relative_path=False, find_stubs_package=find_stubs)
return

name, result = finder.find_spec("package", try_relative_path=False, find_stubs_package=find_stubs)
assert name == "package"

if expect == "both":
assert isinstance(result, Package)
assert result.path.suffix == ".py"
assert not result.path.parent.name.endswith("-stubs")
assert result.stubs
assert result.stubs.suffix == ".pyi"
assert result.stubs.parent.name.endswith("-stubs")
elif expect == "package":
assert isinstance(result, Package)
assert result.path.suffix == ".py"
assert not result.path.parent.name.endswith("-stubs")
assert result.stubs is None
elif expect == "stubs":
assert isinstance(result, Package)
assert result.path.suffix == ".pyi"
assert result.path.parent.name.endswith("-stubs")
assert result.stubs is None
Loading

0 comments on commit 6e55f3b

Please # to comment.