Add support for .pyi files

CHANGELOG.md

@@ -2,10 +2,13 @@
 
 # Unreleased: pdoc next
 
+ - pdoc now picks up type annotations from `.pyi` stub files (PEP-561).
+   This greatly improves support for native modules where no Python source code is available,
+   for example when using PyO3.
+   ([#390](https://github.com/mitmproxy/pdoc/issues/390), [@mhils](https://github.com/mhils))
  - Improve rendering of `typing.TypedDict` subclasses.
    ([#389](https://github.com/mitmproxy/pdoc/issues/389), [@mhils](https://github.com/mhils))
 
-
 # 2022-04-24: pdoc 11.1.0
 
  - Display line numbers when viewing source code.

pdoc/doc.py

@@ -33,7 +33,7 @@
 from pathlib import Path
 from typing import Any, ClassVar, Generic, TypeVar, Union
 
-from pdoc import doc_ast, extract
+from pdoc import doc_ast, doc_pyi, extract
 from pdoc.doc_types import (
     GenericAlias,
     NonUserDefinedCallables,
@@ -304,6 +304,12 @@ def members(self) -> dict[str, Doc]:
             if self._var_docstrings.get(name):
                 doc.docstring = self._var_docstrings[name]
             members[doc.name] = doc
+
+        if isinstance(self, Module):
+            # quirk: doc_pyi expects .members to be set already
+            self.members = members  # type: ignore
+            doc_pyi.include_typeinfo_from_stub_files(self)
+
         return members
 
     @cached_property

pdoc/doc_pyi.py

@@ -0,0 +1,116 @@
+"""
+This module is responsible for patching `pdoc.doc.Doc` objects with type annotations found
+in `.pyi` type stub files ([PEP 561](https://peps.python.org/pep-0561/)).
+This makes it possible to add type hints for native modules such as modules written using [PyO3](https://pyo3.rs/).
+"""
+from __future__ import annotations
+
+from unittest import mock
+
+import sys
+import traceback
+import types
+import warnings
+from pathlib import Path
+
+from pdoc import doc
+from ._compat import cache
+
+
+@cache
+def find_stub_file(module_name: str) -> Path | None:
+    """Try to find a .pyi file with type stubs for the given module name."""
+    module_path = module_name.replace(".", "/")
+
+    for dir in sys.path:
+        file_candidates = [
+            Path(dir) / (module_path + ".pyi"),
+            Path(dir) / (module_path + "/__init__.pyi"),
+        ]
+        for f in file_candidates:
+            if f.exists():
+                return f
+    return None
+
+
+def _import_stub_file(module_name: str, stub_file: Path) -> types.ModuleType:
+    """Import the type stub outside of the normal import machinery."""
+    code = compile(stub_file.read_text(), str(stub_file), "exec")
+    m = types.ModuleType(module_name)
+    m.__file__ = str(stub_file)
+    eval(code, m.__dict__, m.__dict__)
+
+    return m
+
+
+def _prepare_module(ns: doc.Namespace) -> None:
+    """
+    Touch all lazy properties that are accessed in `_patch_doc` to make sure that they are precomputed.
+    We want to do this in advance while sys.modules is not monkeypatched yet.
+    """
+
+    # at the moment, .members is the only lazy property that is accessed.
+    for member in ns.members.values():
+        if isinstance(member, doc.Class):
+            _prepare_module(member)
+
+
+def _patch_doc(target_doc: doc.Doc, stub_mod: doc.Module) -> None:
+    """
+    Patch the target doc (a "real" Python module, e.g. a ".py" file)
+    with the type information from stub_mod (a ".pyi" file).
+    """
+    if target_doc.qualname:
+        stub_doc = stub_mod.get(target_doc.qualname)
+        if stub_doc is None:
+            return
+    else:
+        stub_doc = stub_mod
+
+    if isinstance(target_doc, doc.Function) and isinstance(stub_doc, doc.Function):
+        target_doc.signature = stub_doc.signature
+        target_doc.funcdef = stub_doc.funcdef
+    elif isinstance(target_doc, doc.Variable) and isinstance(stub_doc, doc.Variable):
+        target_doc.annotation = stub_doc.annotation
+    elif isinstance(target_doc, doc.Namespace) and isinstance(stub_doc, doc.Namespace):
+        # pdoc currently does not include variables without docstring in .members (not ideal),
+        # so the regular patching won't work. We manually copy over type annotations instead.
+        for (k, v) in stub_doc._var_annotations.items():
+            var = target_doc.members.get(k, None)
+            if isinstance(var, doc.Variable):
+                var.annotation = v
+
+        for m in target_doc.members.values():
+            _patch_doc(m, stub_mod)
+    else:
+        warnings.warn(
+            f"Error processing type stub for {target_doc.fullname}: "
+            f"Stub is a {stub_doc.type}, but target is a {target_doc.type}."
+        )
+
+
+def include_typeinfo_from_stub_files(module: doc.Module) -> None:
+    # Check if module is a stub module itself - we don't want to recurse!
+    module_file = str(
+        doc._safe_getattr(sys.modules.get(module.modulename), "__file__", "")
+    )
+    if module_file.endswith(".pyi"):
+        return
+
+    stub_file = find_stub_file(module.modulename)
+    if not stub_file:
+        return
+
+    try:
+        imported_stub = _import_stub_file(module.modulename, stub_file)
+    except Exception:
+        warnings.warn(
+            f"Error parsing type stubs for {module.modulename}:\n{traceback.format_exc()}"
+        )
+        return
+
+    _prepare_module(module)
+
+    stub_mod = doc.Module(imported_stub)
+    with mock.patch.dict("sys.modules", {module.modulename: imported_stub}):
+        _patch_doc(module, stub_mod)

test/test_doc_pyi.py

@@ -0,0 +1,32 @@
+import pytest
+import types
+from pathlib import Path
+
+from pdoc import doc, doc_pyi
+
+here = Path(__file__).parent.absolute()
+
+
+def test_type_stub_mismatch():
+    def foo_func():
+        """I'm a func"""
+
+    class foo_cls:
+        """I'm a cls"""
+
+    target = types.ModuleType("test")
+    target.__dict__["foo"] = foo_func
+    target.__dict__["__all__"] = ["foo"]
+
+    stub = types.ModuleType("test")
+    stub.__dict__["foo"] = foo_cls
+    stub.__dict__["__all__"] = ["foo"]
+
+    with pytest.warns(UserWarning, match="Error processing type stub"):
+        doc_pyi._patch_doc(doc.Module(target), doc.Module(stub))
+
+
+def test_invalid_stub_file(monkeypatch):
+    monkeypatch.setattr(doc_pyi, "find_stub_file", lambda m: here / "import_err/err/__init__.py")
+    with pytest.warns(UserWarning, match=r"Error parsing type stubs[\s\S]+RuntimeError"):
+        _ = doc.Module(doc).members

test/test_snapshot.py

@@ -146,6 +146,7 @@ def outfile(self, format: str) -> Path:
     ),
     Snapshot("top_level_reimports", ["top_level_reimports"]),
     Snapshot("type_checking_imports"),
+    Snapshot("type_stub", min_version=(3, 10)),
 ]
 
 

test/testdata/type_stub.py

@@ -0,0 +1,26 @@
+"""
+This module has an accompanying .pyi file with type stubs.
+"""
+
+
+def func(x, y):
+    """A simple function."""
+
+
+var = []
+"""A simple variable."""
+
+
+class Class:
+    attr = 42
+    """An attribute"""
+
+    def meth(self, y):
+        """A simple method."""
+
+    class Subclass:
+        attr = "42"
+        """An attribute"""
+
+        def meth(self, y):
+            """A simple method."""

test/testdata/type_stub.pyi

@@ -0,0 +1,21 @@
+from typing import Any, Iterable
+
+
+def func(x: str, y: Any, z: "Iterable[str]") -> int:
+    ...
+
+
+var: list[str]
+
+
+class Class:
+    attr: int
+
+    def meth(self, y: bool) -> bool:
+        ...
+
+    class Subclass:
+        attr: str
+
+        def meth(self, y: bool) -> bool:
+            ...

test/testdata/type_stub.txt

@@ -0,0 +1,11 @@
+<module type_stub  # This module has an a…
+    <function def func(x: str, y: Any, z: Iterable[str]) -> int: ...  # A simple function.>
+    <var var: list[str] = []  # A simple variable.>
+    <class type_stub.Class
+        <method def __init__(): ...>
+        <var attr: int = 42  # An attribute>
+        <method def meth(self, y: bool) -> bool: ...  # A simple method.>
+        <class type_stub.Class.Subclass
+            <method def __init__(): ...>
+            <var attr: str = '42'  # An attribute>
+            <method def meth(self, y: bool) -> bool: ...  # A simple method.>>>>