feat: Add option to find, load and merge stubs-only packages

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 <[email protected]>

src/griffe/cli.py

@@ -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(
@@ -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:
@@ -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",
@@ -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,
@@ -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.
@@ -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
 
@@ -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,
@@ -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(
@@ -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.

src/griffe/finder.py

@@ -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.
 
@@ -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:
@@ -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.
@@ -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)
@@ -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:

src/griffe/git.py

@@ -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.
 
@@ -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.
@@ -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,
         )
 
 

src/griffe/loader.py

@@ -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.
@@ -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:
@@ -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
 
@@ -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.
 
@@ -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.
@@ -670,6 +687,7 @@ def load(
         module=module,
         submodules=submodules,
         try_relative_path=try_relative_path,
+        find_stubs_package=find_stubs_package,
     )
 
 

src/griffe/merger.py

@@ -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:

tests/test_finder.py

@@ -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
 
 
@@ -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