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

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:
 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)

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