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