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 unittest import mock
  9
 10import sys
 11import traceback
 12import types
 13import warnings
 14from pathlib import Path
 15
 16from pdoc import doc
 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    elif isinstance(target_doc, doc.Variable) and isinstance(stub_doc, doc.Variable):
 74        target_doc.annotation = stub_doc.annotation
 75    elif isinstance(target_doc, doc.Namespace) and isinstance(stub_doc, doc.Namespace):
 76        # pdoc currently does not include variables without docstring in .members (not ideal),
 77        # so the regular patching won't work. We manually copy over type annotations instead.
 78        for (k, v) in stub_doc._var_annotations.items():
 79            var = target_doc.members.get(k, None)
 80            if isinstance(var, doc.Variable):
 81                var.annotation = v
 82
 83        for m in target_doc.members.values():
 84            _patch_doc(m, stub_mod)
 85    else:
 86        warnings.warn(
 87            f"Error processing type stub for {target_doc.fullname}: "
 88            f"Stub is a {stub_doc.type}, but target is a {target_doc.type}."
 89        )
 90
 91
 92def include_typeinfo_from_stub_files(module: doc.Module) -> None:
 93    """Patch the provided module with type information from a matching .pyi file."""
 94    # Check if module is a stub module itself - we don't want to recurse!
 95    module_file = str(
 96        doc._safe_getattr(sys.modules.get(module.modulename), "__file__", "")
 97    )
 98    if module_file.endswith(".pyi"):
 99        return
100
101    stub_file = find_stub_file(module.modulename)
102    if not stub_file:
103        return
104
105    try:
106        imported_stub = _import_stub_file(module.modulename, stub_file)
107    except Exception:
108        warnings.warn(
109            f"Error parsing type stubs for {module.modulename}:\n{traceback.format_exc()}"
110        )
111        return
112
113    _prepare_module(module)
114
115    stub_mod = doc.Module(imported_stub)
116    with mock.patch.dict("sys.modules", {module.modulename: imported_stub}):
117        _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:
 93def include_typeinfo_from_stub_files(module: doc.Module) -> None:
 94    """Patch the provided module with type information from a matching .pyi file."""
 95    # Check if module is a stub module itself - we don't want to recurse!
 96    module_file = str(
 97        doc._safe_getattr(sys.modules.get(module.modulename), "__file__", "")
 98    )
 99    if module_file.endswith(".pyi"):
100        return
101
102    stub_file = find_stub_file(module.modulename)
103    if not stub_file:
104        return
105
106    try:
107        imported_stub = _import_stub_file(module.modulename, stub_file)
108    except Exception:
109        warnings.warn(
110            f"Error parsing type stubs for {module.modulename}:\n{traceback.format_exc()}"
111        )
112        return
113
114    _prepare_module(module)
115
116    stub_mod = doc.Module(imported_stub)
117    with mock.patch.dict("sys.modules", {module.modulename: imported_stub}):
118        _patch_doc(module, stub_mod)

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