Edit on GitHub

pdoc.doc_pyi

This module is responsible for patching pdoc.doc.Doc objects with type annotations found in .pyi type stub files (PEP 561). This makes it possible to add type hints for native modules such as modules written using PyO3.

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

Try to find a .pyi file with type stubs for the given module name.

def include_typeinfo_from_stub_files(module: pdoc.doc.Module) -> None:
 89def include_typeinfo_from_stub_files(module: doc.Module) -> None:
 90    """Patch the provided module with type information from a matching .pyi file."""
 91    # Check if module is a stub module itself - we don't want to recurse!
 92    module_file = str(
 93        doc._safe_getattr(sys.modules.get(module.modulename), "__file__", "")
 94    )
 95    if module_file.endswith(".pyi"):
 96        return
 97
 98    stub_file = find_stub_file(module.modulename)
 99    if not stub_file:
100        return
101
102    try:
103        imported_stub = _import_stub_file(module.modulename, stub_file)
104    except Exception:
105        warnings.warn(
106            f"Error parsing type stubs for {module.modulename}:\n{traceback.format_exc()}"
107        )
108        return
109
110    _prepare_module(module)
111
112    stub_mod = doc.Module(imported_stub)
113    with mock.patch.dict("sys.modules", {module.modulename: imported_stub}):
114        _patch_doc(module, stub_mod)

Patch the provided module with type information from a matching .pyi file.