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