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
  9from pathlib import Path
 10import sys
 11import traceback
 12import types
 13import typing
 14from unittest import mock
 15import warnings
 16
 17from pdoc import doc
 18
 19from ._compat import cache
 20
 21overload_docstr = typing.overload(lambda: None).__doc__
 22
 23
 24@cache
 25def find_stub_file(module_name: str) -> Path | None:
 26    """Try to find a .pyi file with type stubs for the given module name."""
 27    module_path = module_name.replace(".", "/")
 28
 29    # Find .pyi-file in a PEP 0561 compatible stub-package
 30    module_part_name = module_name.split(".")
 31    module_part_name[0] = f"{module_part_name[0]}-stubs"
 32    module_stub_path = "/".join(module_part_name)
 33
 34    for search_dir in sys.path:
 35        file_candidates = [
 36            Path(search_dir) / (module_path + ".pyi"),
 37            Path(search_dir) / (module_path + "/__init__.pyi"),
 38            Path(search_dir) / (module_stub_path + ".pyi"),
 39            Path(search_dir) / (module_stub_path + "/__init__.pyi"),
 40        ]
 41        for f in file_candidates:
 42            if f.exists():
 43                return f
 44    return None
 45
 46
 47def _import_stub_file(module_name: str, stub_file: Path) -> types.ModuleType:
 48    """Import the type stub outside of the normal import machinery."""
 49    code = compile(stub_file.read_text(), str(stub_file), "exec")
 50    m = types.ModuleType(module_name)
 51    m.__file__ = str(stub_file)
 52    eval(code, m.__dict__, m.__dict__)
 53
 54    return m
 55
 56
 57def _prepare_module(ns: doc.Namespace) -> None:
 58    """
 59    Touch all lazy properties that are accessed in `_patch_doc` to make sure that they are precomputed.
 60    We want to do this in advance while sys.modules is not monkeypatched yet.
 61    """
 62
 63    # at the moment, .members is the only lazy property that is accessed.
 64    for member in ns.members.values():
 65        if isinstance(member, doc.Class):
 66            _prepare_module(member)
 67
 68
 69def _patch_doc(target_doc: doc.Doc, stub_mod: doc.Module) -> None:
 70    """
 71    Patch the target doc (a "real" Python module, e.g. a ".py" file)
 72    with the type information from stub_mod (a ".pyi" file).
 73    """
 74    if target_doc.qualname:
 75        stub_doc = stub_mod.get(target_doc.qualname)
 76        if stub_doc is None:
 77            return
 78    else:
 79        stub_doc = stub_mod
 80
 81    if isinstance(target_doc, doc.Function) and isinstance(stub_doc, doc.Function):
 82        # pyi files have functions where all defs have @overload.
 83        # We don't want to pick up the docstring from the typing helper.
 84        if stub_doc.docstring == overload_docstr:
 85            stub_doc.docstring = ""
 86
 87        target_doc.signature = stub_doc.signature
 88        target_doc.funcdef = stub_doc.funcdef
 89        target_doc.docstring = stub_doc.docstring or target_doc.docstring
 90    elif isinstance(target_doc, doc.Variable) and isinstance(stub_doc, doc.Variable):
 91        target_doc.annotation = stub_doc.annotation
 92        target_doc.docstring = stub_doc.docstring or target_doc.docstring
 93    elif isinstance(target_doc, doc.Namespace) and isinstance(stub_doc, doc.Namespace):
 94        target_doc.docstring = stub_doc.docstring or target_doc.docstring
 95        for m in target_doc.members.values():
 96            _patch_doc(m, stub_mod)
 97    else:
 98        warnings.warn(
 99            f"Error processing type stub for {target_doc.fullname}: "
100            f"Stub is a {stub_doc.kind}, but target is a {target_doc.kind}."
101        )
102
103
104def include_typeinfo_from_stub_files(module: doc.Module) -> None:
105    """Patch the provided module with type information from a matching .pyi file."""
106    # Check if module is a stub module itself - we don't want to recurse!
107    module_file = str(
108        doc._safe_getattr(sys.modules.get(module.modulename), "__file__", "")
109    )
110    if module_file.endswith(".pyi"):
111        return
112
113    stub_file = find_stub_file(module.modulename)
114    if not stub_file:
115        return
116
117    try:
118        imported_stub = _import_stub_file(module.modulename, stub_file)
119    except Exception:
120        warnings.warn(
121            f"Error parsing type stubs for {module.modulename}:\n{traceback.format_exc()}"
122        )
123        return
124
125    _prepare_module(module)
126
127    stub_mod = doc.Module(imported_stub)
128    with mock.patch.dict("sys.modules", {module.modulename: imported_stub}):
129        _patch_doc(module, stub_mod)