feat: Add option to give precedence to docstrings in stubs vs. source

packages/griffelib/src/griffe/_internal/collections.py

@@ -68,6 +68,10 @@ class ModulesCollection(GetMembersMixin, SetMembersMixin, DelMembersMixin):
     is_collection = True
     """Marked as collection to distinguish from objects."""
 
+    # "Prefer stubs docstrings": we store it here
+    # to be able to access it in `SetMembersMixin.set_member`.
+    _psd: bool = False
+
     def __init__(self) -> None:
         """Initialize the collection."""
         self.members: dict[str, Module] = {}

packages/griffelib/src/griffe/_internal/helpers.py

@@ -135,6 +135,7 @@ def temporary_visited_package(
     resolve_external: bool | None = None,
     resolve_implicit: bool = False,
     search_sys_path: bool = False,
+    prefer_stubs_docs: bool = False,
 ) -> Iterator[Module]:
     """Create and visit a temporary package.
 
@@ -162,6 +163,9 @@ def temporary_visited_package(
             or the origin module (for example when `ast` imports from `_ast`).
         resolve_implicit: When false, only try to resolve an alias if it is explicitly exported.
         search_sys_path: Whether to search the system paths for the package.
+        prefer_stubs_docs: Whether to give precedence to stubs docstrings
+            rather than source docstrings. When both are present, the stubs one
+            will override the source one.
 
     Yields:
         A module.
@@ -182,6 +186,7 @@ def temporary_visited_package(
             resolve_external=resolve_external,
             resolve_implicit=resolve_implicit,
             force_inspection=False,
+            prefer_stubs_docs=prefer_stubs_docs,
         )
 
 
@@ -203,6 +208,7 @@ def temporary_inspected_package(
     resolve_external: bool | None = None,
     resolve_implicit: bool = False,
     search_sys_path: bool = False,
+    prefer_stubs_docs: bool = False,
 ) -> Iterator[Module]:
     """Create and inspect a temporary package.
 
@@ -230,6 +236,9 @@ def temporary_inspected_package(
             or the origin module (for example when `ast` imports from `_ast`).
         resolve_implicit: When false, only try to resolve an alias if it is explicitly exported.
         search_sys_path: Whether to search the system paths for the package.
+        prefer_stubs_docs: Whether to give precedence to stubs docstrings
+            rather than source docstrings. When both are present, the stubs one
+            will override the source one.
 
     Yields:
         A module.
@@ -251,6 +260,7 @@ def temporary_inspected_package(
                 resolve_external=resolve_external,
                 resolve_implicit=resolve_implicit,
                 force_inspection=True,
+                prefer_stubs_docs=prefer_stubs_docs,
             )
         finally:
             for name in tuple(sys.modules.keys()):

packages/griffelib/src/griffe/_internal/loader.py

@@ -62,6 +62,7 @@ def __init__(
         allow_inspection: bool = True,
         force_inspection: bool = False,
         store_source: bool = True,
+        prefer_stubs_docs: bool = False,
     ) -> None:
         """Initialize the loader.
 
@@ -74,6 +75,9 @@ def __init__(
             modules_collection: A collection of modules.
             allow_inspection: Whether to allow inspecting modules when visiting them is not possible.
             store_source: Whether to store code source in the lines collection.
+            prefer_stubs_docs: Whether to give precedence to stubs docstrings
+                rather than source docstrings. When both are present, the stubs one
+                will override the source one.
         """
         self.extensions: Extensions = extensions or load_extensions()
         """Loaded Griffe extensions."""
@@ -84,13 +88,16 @@ def __init__(
         self.lines_collection: LinesCollection = lines_collection or LinesCollection()
         """Collection of source code lines."""
         self.modules_collection: ModulesCollection = modules_collection or ModulesCollection()
+        self.modules_collection._psd = prefer_stubs_docs
         """Collection of modules."""
         self.allow_inspection: bool = allow_inspection
         """Whether to allow inspecting (importing) modules for which we can't find sources."""
         self.force_inspection: bool = force_inspection
         """Whether to force inspecting (importing) modules, even when sources were found."""
         self.store_source: bool = store_source
         """Whether to store source code in the lines collection."""
+        self.prefer_stubs_docs: bool = prefer_stubs_docs
+        """Whether to give precedence to stubs docstrings rather than source ones."""
         self._search_paths: Sequence[str | Path] | None = search_paths
         self._time_stats: dict = {
             "time_spent_visiting": 0,
@@ -548,7 +555,7 @@ def _load_package(self, package: Package | NamespacePackage, *, submodules: bool
             # 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 merge_stubs(top_module, stubs, prefer_stubs_docs=self.prefer_stubs_docs)
         return top_module
 
     def _load_module(
@@ -760,6 +767,7 @@ def load(
     force_inspection: bool = False,
     store_source: bool = True,
     find_stubs_package: bool = False,
+    prefer_stubs_docs: bool = False,
     resolve_aliases: bool = False,
     resolve_external: bool | None = None,
     resolve_implicit: bool = False,
@@ -818,6 +826,9 @@ def load(
         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.
+        prefer_stubs_docs: Whether to give precedence to stubs docstrings
+            rather than source docstrings. When both are present, the stubs one
+            will override the source one.
         resolve_aliases: Whether to resolve aliases.
         resolve_external: Whether to try to load unspecified modules to resolve aliases.
             Default value (`None`) means to load external modules only if they are the private sibling
@@ -837,6 +848,7 @@ def load(
         allow_inspection=allow_inspection,
         force_inspection=force_inspection,
         store_source=store_source,
+        prefer_stubs_docs=prefer_stubs_docs,
     )
     result = loader.load(
         objspec,
@@ -865,6 +877,7 @@ def load_git(
     allow_inspection: bool = True,
     force_inspection: bool = False,
     find_stubs_package: bool = False,
+    prefer_stubs_docs: bool = False,
     resolve_aliases: bool = False,
     resolve_external: bool | None = None,
     resolve_implicit: bool = False,
@@ -901,6 +914,9 @@ def load_git(
         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.
+        prefer_stubs_docs: Whether to give precedence to stubs docstrings
+            rather than source docstrings. When both are present, the stubs one
+            will override the source one.
         resolve_aliases: Whether to resolve aliases.
         resolve_external: Whether to try to load unspecified modules to resolve aliases.
             Default value (`None`) means to load external modules only if they are the private sibling
@@ -931,6 +947,7 @@ def load_git(
             resolve_aliases=resolve_aliases,
             resolve_external=resolve_external,
             resolve_implicit=resolve_implicit,
+            prefer_stubs_docs=prefer_stubs_docs,
         )
 
 
@@ -949,6 +966,7 @@ def load_pypi(
     allow_inspection: bool = True,
     force_inspection: bool = False,
     find_stubs_package: bool = False,
+    prefer_stubs_docs: bool = False,
     resolve_aliases: bool = False,
     resolve_external: bool | None = None,
     resolve_implicit: bool = False,
@@ -972,6 +990,9 @@ def load_pypi(
         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.
+        prefer_stubs_docs: Whether to give precedence to stubs docstrings
+            rather than source docstrings. When both are present, the stubs one
+            will override the source one.
         resolve_aliases: Whether to resolve aliases.
         resolve_external: Whether to try to load unspecified modules to resolve aliases.
             Default value (`None`) means to load external modules only if they are the private sibling
@@ -1056,4 +1077,5 @@ def load_pypi(
         resolve_aliases=resolve_aliases,
         resolve_external=resolve_external,
         resolve_implicit=resolve_implicit,
+        prefer_stubs_docs=prefer_stubs_docs,
     )

packages/griffelib/src/griffe/_internal/merger.py

@@ -15,42 +15,42 @@
     from griffe._internal.models import Attribute, Class, Function, Module, Object, TypeAlias
 
 
-def _merge_module_stubs(module: Module, stubs: Module) -> None:
-    _merge_stubs_docstring(module, stubs)
+def _merge_module_stubs(module: Module, stubs: Module, *, psd: bool = False) -> None:
+    _merge_stubs_docstring(module, stubs, psd=psd)
     _merge_stubs_overloads(module, stubs)
-    _merge_stubs_members(module, stubs)
+    _merge_stubs_members(module, stubs, psd=psd)
 
 
-def _merge_class_stubs(class_: Class, stubs: Class) -> None:
-    _merge_stubs_docstring(class_, stubs)
+def _merge_class_stubs(class_: Class, stubs: Class, *, psd: bool = False) -> None:
+    _merge_stubs_docstring(class_, stubs, psd=psd)
     _merge_stubs_overloads(class_, stubs)
     _merge_stubs_type_parameters(class_, stubs)
-    _merge_stubs_members(class_, stubs)
+    _merge_stubs_members(class_, stubs, psd=psd)
 
 
-def _merge_function_stubs(function: Function, stubs: Function) -> None:
-    _merge_stubs_docstring(function, stubs)
+def _merge_function_stubs(function: Function, stubs: Function, *, psd: bool = False) -> None:
+    _merge_stubs_docstring(function, stubs, psd=psd)
     for parameter in stubs.parameters:
         with suppress(KeyError):
             function.parameters[parameter.name].annotation = parameter.annotation
     function.returns = stubs.returns
     _merge_stubs_type_parameters(function, stubs)
 
 
-def _merge_attribute_stubs(attribute: Attribute, stubs: Attribute) -> None:
-    _merge_stubs_docstring(attribute, stubs)
+def _merge_attribute_stubs(attribute: Attribute, stubs: Attribute, *, psd: bool = False) -> None:
+    _merge_stubs_docstring(attribute, stubs, psd=psd)
     attribute.annotation = stubs.annotation
     if stubs.value not in (None, "..."):
         attribute.value = stubs.value
 
 
-def _merge_type_alias_stubs(type_alias: TypeAlias, stubs: TypeAlias) -> None:
-    _merge_stubs_docstring(type_alias, stubs)
+def _merge_type_alias_stubs(type_alias: TypeAlias, stubs: TypeAlias, *, psd: bool = False) -> None:
+    _merge_stubs_docstring(type_alias, stubs, psd=psd)
     _merge_stubs_type_parameters(type_alias, stubs)
 
 
-def _merge_stubs_docstring(obj: Object, stubs: Object) -> None:
-    if not obj.docstring and stubs.docstring:
+def _merge_stubs_docstring(obj: Object, stubs: Object, *, psd: bool = False) -> None:
+    if (psd or not obj.docstring) and stubs.docstring:
         obj.docstring = stubs.docstring
 
 
@@ -107,7 +107,7 @@ def _merge_overload_annotations(function: Function, overloads: list[Function]) -
         function.returns = _merge_annotations(return_annotations)
 
 
-def _merge_stubs_members(obj: Module | Class, stubs: Module | Class) -> None:
+def _merge_stubs_members(obj: Module | Class, stubs: Module | Class, *, psd: bool = False) -> None:
     # Merge imports to later know if objects coming from the stubs were imported.
     obj.imports.update(stubs.imports)
 
@@ -137,26 +137,29 @@ def _merge_stubs_members(obj: Module | Class, stubs: Module | Class) -> None:
                     )
                     obj.set_member(stub_member.name, stub_member)
                 elif obj_member.is_module:
-                    _merge_module_stubs(obj_member, stub_member)  # ty:ignore[invalid-argument-type]
+                    _merge_module_stubs(obj_member, stub_member, psd=psd)  # ty:ignore[invalid-argument-type]
                 elif obj_member.is_class:
-                    _merge_class_stubs(obj_member, stub_member)  # ty:ignore[invalid-argument-type]
+                    _merge_class_stubs(obj_member, stub_member, psd=psd)  # ty:ignore[invalid-argument-type]
                 elif obj_member.is_function:
-                    _merge_function_stubs(obj_member, stub_member)  # ty:ignore[invalid-argument-type]
+                    _merge_function_stubs(obj_member, stub_member, psd=psd)  # ty:ignore[invalid-argument-type]
                 elif obj_member.is_attribute:
-                    _merge_attribute_stubs(obj_member, stub_member)  # ty:ignore[invalid-argument-type]
+                    _merge_attribute_stubs(obj_member, stub_member, psd=psd)  # ty:ignore[invalid-argument-type]
                 elif obj_member.is_type_alias:
-                    _merge_type_alias_stubs(obj_member, stub_member)  # ty:ignore[invalid-argument-type]
+                    _merge_type_alias_stubs(obj_member, stub_member, psd=psd)  # ty:ignore[invalid-argument-type]
         else:
             stub_member.runtime = False
             obj.set_member(member_name, stub_member)
 
 
-def merge_stubs(mod1: Module, mod2: Module) -> Module:
+def merge_stubs(mod1: Module, mod2: Module, *, prefer_stubs_docs: bool = False) -> Module:
     """Merge stubs into a module.
 
     Parameters:
         mod1: A regular module or stubs module.
         mod2: A regular module or stubs module.
+        prefer_stubs_docs: Whether to give precedence to stubs docstrings
+            rather than source docstrings. When both are present, the stubs one
+            will override the source one.
 
     Raises:
         ValueError: When both modules are regular modules (no stubs is passed).
@@ -173,5 +176,5 @@ def merge_stubs(mod1: Module, mod2: Module) -> Module:
         module = mod1
     else:
         raise ValueError("cannot merge regular (non-stubs) modules together")
-    _merge_module_stubs(module, stubs)
+    _merge_module_stubs(module, stubs, psd=prefer_stubs_docs)
     return module

packages/griffelib/src/griffe/_internal/mixins.py

@@ -195,8 +195,9 @@ def set_member(self, key: str | Sequence[str], value: Object | Alias) -> None:
                         # Accessing file paths can trigger a builtin module error.
                         with suppress(AliasResolutionError, CyclicAliasError, BuiltinModuleError):
                             if value.is_module and value.filepath != member.filepath:
+                                psd = self._psd if self.is_collection else self.modules_collection._psd  # ty:ignore[unresolved-attribute]
                                 with suppress(ValueError):
-                                    value = merge_stubs(member, value)  # ty:ignore[invalid-argument-type]
+                                    value = merge_stubs(member, value, prefer_stubs_docs=psd)  # ty:ignore[invalid-argument-type]
                     for alias in member.aliases.values():
                         with suppress(CyclicAliasError):
                             alias.target = value

packages/griffelib/tests/test_merger.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import pytest
+
 from griffe import temporary_visited_package
 
 
@@ -95,3 +97,21 @@ def func(x: float) -> float: ...
         func = pkg["mod.func"]
         assert str(func.parameters["x"].annotation) == "int | float"
         assert str(func.returns) == "int | float"
+
+
+@pytest.mark.parametrize(
+    ("psd", "expected"),
+    [(True, "Stubs"), (False, "Source")],
+)
+def test_prefer_stubs_docstrings(psd: bool, expected: str) -> None:
+    """The "prefer stubs docstrings" option is respected."""
+    with temporary_visited_package(
+        "package",
+        {
+            "mod.py": "def func():\n    '''Source'''",
+            "mod.pyi": "def func():\n    '''Stubs'''",
+        },
+        prefer_stubs_docs=psd,
+    ) as pkg:
+        func = pkg["mod.func"]
+        assert func.docstring.value == expected