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"""
  6
  7from __future__ import annotations
  8
  9import importlib.util
 10from pathlib import Path
 11import sys
 12import traceback
 13import types
 14import typing
 15from unittest import mock
 16import warnings
 17
 18from pdoc import doc
 19
 20from ._compat import cache
 21
 22overload_docstr = typing.overload(lambda: None).__doc__
 23
 24
 25@cache
 26def find_stub_file(module_name: str) -> Path | None:
 27    """Try to find a .pyi file with type stubs for the given module name."""
 28    module_path = module_name.replace(".", "/")
 29
 30    # Find .pyi-file in a PEP 0561 compatible stub-package
 31    module_part_name = module_name.split(".")
 32    module_part_name[0] = f"{module_part_name[0]}-stubs"
 33    module_stub_path = "/".join(module_part_name)
 34
 35    for search_dir in sys.path:
 36        file_candidates = [
 37            Path(search_dir) / (module_path + ".pyi"),
 38            Path(search_dir) / (module_path + "/__init__.pyi"),
 39            Path(search_dir) / (module_stub_path + ".pyi"),
 40            Path(search_dir) / (module_stub_path + "/__init__.pyi"),
 41        ]
 42        for f in file_candidates:
 43            if f.exists():
 44                return f
 45    return None
 46
 47
 48def _import_stub_file(module_name: str, stub_file: Path) -> types.ModuleType:
 49    """
 50    Import the type stub outside of the normal import machinery.
 51
 52    Note that currently, for objects imported by the stub file, the _original_ module
 53    is used and not the corresponding stub file.
 54    """
 55    sys.path_hooks.append(
 56        importlib.machinery.FileFinder.path_hook(
 57            (importlib.machinery.SourceFileLoader, [".pyi"])
 58        )
 59    )
 60    try:
 61        loader = importlib.machinery.SourceFileLoader(module_name, str(stub_file))
 62        spec = importlib.util.spec_from_file_location(
 63            module_name, stub_file, loader=loader
 64        )
 65        assert spec is not None
 66        m = importlib.util.module_from_spec(spec)
 67        loader.exec_module(m)
 68        return m
 69    finally:
 70        sys.path_hooks.pop()
 71
 72
 73def _prepare_module(ns: doc.Namespace) -> None:
 74    """
 75    Touch all lazy properties that are accessed in `_patch_doc` to make sure that they are precomputed.
 76    We want to do this in advance while sys.modules is not monkeypatched yet.
 77    """
 78
 79    # at the moment, .members is the only lazy property that is accessed.
 80    for member in ns.members.values():
 81        if isinstance(member, doc.Class):
 82            _prepare_module(member)
 83
 84
 85def _patch_doc(target_doc: doc.Doc, stub_mod: doc.Module) -> None:
 86    """
 87    Patch the target doc (a "real" Python module, e.g. a ".py" file)
 88    with the type information from stub_mod (a ".pyi" file).
 89    """
 90    if target_doc.qualname:
 91        stub_doc = stub_mod.get(target_doc.qualname)
 92        if stub_doc is None:
 93            return
 94    else:
 95        stub_doc = stub_mod
 96
 97    if isinstance(target_doc, doc.Function) and isinstance(stub_doc, doc.Function):
 98        # pyi files have functions where all defs have @overload.
 99        # We don't want to pick up the docstring from the typing helper.
100        if stub_doc.docstring == overload_docstr:
101            stub_doc.docstring = ""
102
103        target_doc.signature = stub_doc.signature
104        target_doc.funcdef = stub_doc.funcdef
105        target_doc.docstring = stub_doc.docstring or target_doc.docstring
106    elif isinstance(target_doc, doc.Variable) and isinstance(stub_doc, doc.Variable):
107        target_doc.annotation = stub_doc.annotation
108        target_doc.docstring = stub_doc.docstring or target_doc.docstring
109    elif isinstance(target_doc, doc.Namespace) and isinstance(stub_doc, doc.Namespace):
110        target_doc.docstring = stub_doc.docstring or target_doc.docstring
111        for m in target_doc.members.values():
112            _patch_doc(m, stub_mod)
113    else:
114        warnings.warn(
115            f"Error processing type stub for {target_doc.fullname}: "
116            f"Stub is a {stub_doc.kind}, but target is a {target_doc.kind}."
117        )
118
119
120def include_typeinfo_from_stub_files(module: doc.Module) -> None:
121    """Patch the provided module with type information from a matching .pyi file."""
122    # Check if module is a stub module itself - we don't want to recurse!
123    module_file = str(
124        doc._safe_getattr(sys.modules.get(module.modulename), "__file__", "")
125    )
126    if module_file.endswith(".pyi"):
127        return
128
129    stub_file = find_stub_file(module.modulename)
130    if not stub_file:
131        return
132
133    try:
134        imported_stub = _import_stub_file(module.modulename, stub_file)
135    except Exception:
136        warnings.warn(
137            f"Error parsing type stubs for {module.modulename}:\n{traceback.format_exc()}"
138        )
139        return
140
141    _prepare_module(module)
142
143    stub_mod = doc.Module(imported_stub)
144    with mock.patch.dict("sys.modules", {module.modulename: imported_stub}):
145        _patch_doc(module, stub_mod)
overload_docstr = 'Helper for @overload to raise when called.'
@cache
def find_stub_file(module_name: str) -> pathlib.Path | None:
26@cache
27def find_stub_file(module_name: str) -> Path | None:
28    """Try to find a .pyi file with type stubs for the given module name."""
29    module_path = module_name.replace(".", "/")
30
31    # Find .pyi-file in a PEP 0561 compatible stub-package
32    module_part_name = module_name.split(".")
33    module_part_name[0] = f"{module_part_name[0]}-stubs"
34    module_stub_path = "/".join(module_part_name)
35
36    for search_dir in sys.path:
37        file_candidates = [
38            Path(search_dir) / (module_path + ".pyi"),
39            Path(search_dir) / (module_path + "/__init__.pyi"),
40            Path(search_dir) / (module_stub_path + ".pyi"),
41            Path(search_dir) / (module_stub_path + "/__init__.pyi"),
42        ]
43        for f in file_candidates:
44            if f.exists():
45                return f
46    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:
121def include_typeinfo_from_stub_files(module: doc.Module) -> None:
122    """Patch the provided module with type information from a matching .pyi file."""
123    # Check if module is a stub module itself - we don't want to recurse!
124    module_file = str(
125        doc._safe_getattr(sys.modules.get(module.modulename), "__file__", "")
126    )
127    if module_file.endswith(".pyi"):
128        return
129
130    stub_file = find_stub_file(module.modulename)
131    if not stub_file:
132        return
133
134    try:
135        imported_stub = _import_stub_file(module.modulename, stub_file)
136    except Exception:
137        warnings.warn(
138            f"Error parsing type stubs for {module.modulename}:\n{traceback.format_exc()}"
139        )
140        return
141
142    _prepare_module(module)
143
144    stub_mod = doc.Module(imported_stub)
145    with mock.patch.dict("sys.modules", {module.modulename: imported_stub}):
146        _patch_doc(module, stub_mod)

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