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 functools import cache
 10import importlib.util
 11from pathlib import Path
 12import sys
 13import traceback
 14import types
 15import typing
 16from unittest import mock
 17import warnings
 18
 19from pdoc import doc
 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    """
 49    Import the type stub outside of the normal import machinery.
 50
 51    Note that currently, for objects imported by the stub file, the _original_ module
 52    is used and not the corresponding stub file.
 53    """
 54    sys.path_hooks.append(
 55        importlib.machinery.FileFinder.path_hook(
 56            (importlib.machinery.SourceFileLoader, [".pyi"])
 57        )
 58    )
 59    try:
 60        loader = importlib.machinery.SourceFileLoader(module_name, str(stub_file))
 61        spec = importlib.util.spec_from_file_location(
 62            module_name, stub_file, loader=loader
 63        )
 64        assert spec is not None
 65        m = importlib.util.module_from_spec(spec)
 66        loader.exec_module(m)
 67        return m
 68    finally:
 69        sys.path_hooks.pop()
 70
 71
 72def _prepare_module(ns: doc.Namespace) -> None:
 73    """
 74    Touch all lazy properties that are accessed in `_patch_doc` to make sure that they are precomputed.
 75    We want to do this in advance while sys.modules is not monkeypatched yet.
 76    """
 77
 78    # at the moment, .members is the only lazy property that is accessed.
 79    for member in ns.members.values():
 80        if isinstance(member, doc.Class):
 81            _prepare_module(member)
 82
 83
 84def _patch_doc(target_doc: doc.Doc, stub_mod: doc.Module) -> None:
 85    """
 86    Patch the target doc (a "real" Python module, e.g. a ".py" file)
 87    with the type information from stub_mod (a ".pyi" file).
 88    """
 89    if target_doc.qualname:
 90        stub_doc = stub_mod.get(target_doc.qualname)
 91        if stub_doc is None:
 92            return
 93    else:
 94        stub_doc = stub_mod
 95
 96    if isinstance(target_doc, doc.Function) and isinstance(stub_doc, doc.Function):
 97        # pyi files have functions where all defs have @overload.
 98        # We don't want to pick up the docstring from the typing helper.
 99        if stub_doc.docstring == overload_docstr:
100            stub_doc.docstring = ""
101
102        target_doc.signature = stub_doc.signature
103        target_doc.funcdef = stub_doc.funcdef
104        target_doc.docstring = stub_doc.docstring or target_doc.docstring
105    elif isinstance(target_doc, doc.Variable) and isinstance(stub_doc, doc.Variable):
106        target_doc.annotation = stub_doc.annotation
107        target_doc.docstring = stub_doc.docstring or target_doc.docstring
108    elif isinstance(target_doc, doc.Namespace) and isinstance(stub_doc, doc.Namespace):
109        target_doc.docstring = stub_doc.docstring or target_doc.docstring
110        for m in target_doc.members.values():
111            _patch_doc(m, stub_mod)
112    else:
113        warnings.warn(
114            f"Error processing type stub for {target_doc.fullname}: "
115            f"Stub is a {stub_doc.kind}, but target is a {target_doc.kind}."
116        )
117
118
119def include_typeinfo_from_stub_files(module: doc.Module) -> None:
120    """Patch the provided module with type information from a matching .pyi file."""
121    # Check if module is a stub module itself - we don't want to recurse!
122    module_file = str(
123        doc._safe_getattr(sys.modules.get(module.modulename), "__file__", "")
124    )
125    if module_file.endswith(".pyi"):
126        return
127
128    stub_file = find_stub_file(module.modulename)
129    if not stub_file:
130        return
131
132    try:
133        imported_stub = _import_stub_file(module.modulename, stub_file)
134    except Exception:
135        warnings.warn(
136            f"Error parsing type stubs for {module.modulename}:\n{traceback.format_exc()}"
137        )
138        return
139
140    _prepare_module(module)
141
142    stub_mod = doc.Module(imported_stub)
143    with mock.patch.dict("sys.modules", {module.modulename: imported_stub}):
144        _patch_doc(module, stub_mod)