pdoc.doc_pyi

  1"""
  2This module is responsible for patching `pdoc.doc.Doc` objects with type annotations found
  3in `.pyi` type stub files ([PEP 561](https://peps.python.org/pep-0561/)).
  4This makes it possible to add type hints for native modules such as modules written using [PyO3](https://pyo3.rs/).
  5"""
  6
  7from __future__ import annotations
  8
  9import importlib.util
 10from pathlib import Path
 11import sys
 12import traceback
 13import types
 14import typing
 15from unittest import mock
 16import warnings
 17
 18from pdoc import doc
 19
 20from ._compat import cache
 21
 22overload_docstr = typing.overload(lambda: None).__doc__
 23
 24
 25@cache
 26def find_stub_file(module_name: str) -> Path | None:
 27    """Try to find a .pyi file with type stubs for the given module name."""
 28    module_path = module_name.replace(".", "/")
 29
 30    # Find .pyi-file in a PEP 0561 compatible stub-package
 31    module_part_name = module_name.split(".")
 32    module_part_name[0] = f"{module_part_name[0]}-stubs"
 33    module_stub_path = "/".join(module_part_name)
 34
 35    for search_dir in sys.path:
 36        file_candidates = [
 37            Path(search_dir) / (module_path + ".pyi"),
 38            Path(search_dir) / (module_path + "/__init__.pyi"),
 39            Path(search_dir) / (module_stub_path + ".pyi"),
 40            Path(search_dir) / (module_stub_path + "/__init__.pyi"),
 41        ]
 42        for f in file_candidates:
 43            if f.exists():
 44                return f
 45    return None
 46
 47
 48def _import_stub_file(module_name: str, stub_file: Path) -> types.ModuleType:
 49    """
 50    Import the type stub outside of the normal import machinery.
 51
 52    Note that currently, for objects imported by the stub file, the _original_ module
 53    is used and not the corresponding stub file.
 54    """
 55    sys.path_hooks.append(
 56        importlib.machinery.FileFinder.path_hook(
 57            (importlib.machinery.SourceFileLoader, [".pyi"])
 58        )
 59    )
 60    try:
 61        loader = importlib.machinery.SourceFileLoader(module_name, str(stub_file))
 62        spec = importlib.util.spec_from_file_location(
 63            module_name, stub_file, loader=loader
 64        )
 65        assert spec is not None
 66        m = importlib.util.module_from_spec(spec)
 67        loader.exec_module(m)
 68        return m
 69    finally:
 70        sys.path_hooks.pop()
 71
 72
 73def _prepare_module(ns: doc.Namespace) -> None:
 74    """
 75    Touch all lazy properties that are accessed in `_patch_doc` to make sure that they are precomputed.
 76    We want to do this in advance while sys.modules is not monkeypatched yet.
 77    """
 78
 79    # at the moment, .members is the only lazy property that is accessed.
 80    for member in ns.members.values():
 81        if isinstance(member, doc.Class):
 82            _prepare_module(member)
 83
 84
 85def _patch_doc(target_doc: doc.Doc, stub_mod: doc.Module) -> None:
 86    """
 87    Patch the target doc (a "real" Python module, e.g. a ".py" file)
 88    with the type information from stub_mod (a ".pyi" file).
 89    """
 90    if target_doc.qualname:
 91        stub_doc = stub_mod.get(target_doc.qualname)
 92        if stub_doc is None:
 93            return
 94    else:
 95        stub_doc = stub_mod
 96
 97    if isinstance(target_doc, doc.Function) and isinstance(stub_doc, doc.Function):
 98        # pyi files have functions where all defs have @overload.
 99        # We don't want to pick up the docstring from the typing helper.
100        if stub_doc.docstring == overload_docstr:
101            stub_doc.docstring = ""
102
103        target_doc.signature = stub_doc.signature
104        target_doc.funcdef = stub_doc.funcdef
105        target_doc.docstring = stub_doc.docstring or target_doc.docstring
106    elif isinstance(target_doc, doc.Variable) and isinstance(stub_doc, doc.Variable):
107        target_doc.annotation = stub_doc.annotation
108        target_doc.docstring = stub_doc.docstring or target_doc.docstring
109    elif isinstance(target_doc, doc.Namespace) and isinstance(stub_doc, doc.Namespace):
110        target_doc.docstring = stub_doc.docstring or target_doc.docstring
111        for m in target_doc.members.values():
112            _patch_doc(m, stub_mod)
113    else:
114        warnings.warn(
115            f"Error processing type stub for {target_doc.fullname}: "
116            f"Stub is a {stub_doc.kind}, but target is a {target_doc.kind}."
117        )
118
119
120def include_typeinfo_from_stub_files(module: doc.Module) -> None:
121    """Patch the provided module with type information from a matching .pyi file."""
122    # Check if module is a stub module itself - we don't want to recurse!
123    module_file = str(
124        doc._safe_getattr(sys.modules.get(module.modulename), "__file__", "")
125    )
126    if module_file.endswith(".pyi"):
127        return
128
129    stub_file = find_stub_file(module.modulename)
130    if not stub_file:
131        return
132
133    try:
134        imported_stub = _import_stub_file(module.modulename, stub_file)
135    except Exception:
136        warnings.warn(
137            f"Error parsing type stubs for {module.modulename}:\n{traceback.format_exc()}"
138        )
139        return
140
141    _prepare_module(module)
142
143    stub_mod = doc.Module(imported_stub)
144    with mock.patch.dict("sys.modules", {module.modulename: imported_stub}):
145        _patch_doc(module, stub_mod)