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