diff --git a/packages/griffelib/src/griffe/_internal/collections.py b/packages/griffelib/src/griffe/_internal/collections.py index 787f57a7..0bfb26dd 100644 --- a/packages/griffelib/src/griffe/_internal/collections.py +++ b/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] = {} diff --git a/packages/griffelib/src/griffe/_internal/helpers.py b/packages/griffelib/src/griffe/_internal/helpers.py index 46ea94d6..6f57a503 100644 --- a/packages/griffelib/src/griffe/_internal/helpers.py +++ b/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()): diff --git a/packages/griffelib/src/griffe/_internal/loader.py b/packages/griffelib/src/griffe/_internal/loader.py index a9c4edcd..10ef9451 100644 --- a/packages/griffelib/src/griffe/_internal/loader.py +++ b/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,6 +88,7 @@ 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.""" @@ -91,6 +96,8 @@ def __init__( """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, ) diff --git a/packages/griffelib/src/griffe/_internal/merger.py b/packages/griffelib/src/griffe/_internal/merger.py index 3263a63a..99b3b73b 100644 --- a/packages/griffelib/src/griffe/_internal/merger.py +++ b/packages/griffelib/src/griffe/_internal/merger.py @@ -15,21 +15,21 @@ 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 @@ -37,20 +37,20 @@ def _merge_function_stubs(function: Function, stubs: Function) -> None: _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 diff --git a/packages/griffelib/src/griffe/_internal/mixins.py b/packages/griffelib/src/griffe/_internal/mixins.py index 74eb0c6a..057dc8c3 100644 --- a/packages/griffelib/src/griffe/_internal/mixins.py +++ b/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 diff --git a/packages/griffelib/tests/test_merger.py b/packages/griffelib/tests/test_merger.py index c700238e..a1c7e1a6 100644 --- a/packages/griffelib/tests/test_merger.py +++ b/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