improve pyi file import (#732)

CHANGELOG.md

@@ -7,6 +7,8 @@
 - Fix a bug where entire modules would be excluded by `--no-include-undocumented`.
   To exclude modules, see https://pdoc.dev/docs/pdoc.html#exclude-submodules-from-being-documented.
   ([#728](https://github.com/mitmproxy/pdoc/pull/728), @mhils)
+- Fix a bug where pdoc would crash when importing pyi files.
+  ([#732](https://github.com/mitmproxy/pdoc/pull/732), @mhils)
 - Fix a bug where subclasses of TypedDict subclasses would not render correctly.
   ([#729](https://github.com/mitmproxy/pdoc/pull/729), @mhils)
 

pdoc/doc_pyi.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import importlib.util
 from pathlib import Path
 import sys
 import traceback
@@ -45,13 +46,28 @@ def find_stub_file(module_name: str) -> Path | 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__)
+    """
+    Import the type stub outside of the normal import machinery.
 
-    return m
+    Note that currently, for objects imported by the stub file, the _original_ module
+    is used and not the corresponding stub file.
+    """
+    sys.path_hooks.append(
+        importlib.machinery.FileFinder.path_hook(
+            (importlib.machinery.SourceFileLoader, [".pyi"])
+        )
+    )
+    try:
+        loader = importlib.machinery.SourceFileLoader(module_name, str(stub_file))
+        spec = importlib.util.spec_from_file_location(
+            module_name, stub_file, loader=loader
+        )
+        assert spec is not None
+        m = importlib.util.module_from_spec(spec)
+        loader.exec_module(m)
+        return m
+    finally:
+        sys.path_hooks.pop()
 
 
 def _prepare_module(ns: doc.Namespace) -> None:

test/test_snapshot.py

@@ -164,7 +164,7 @@ def outfile(self, format: str) -> Path:
     Snapshot("pyo3_sample_library", specs=["pdoc_pyo3_sample_library"]),
     Snapshot("top_level_reimports", ["top_level_reimports"]),
     Snapshot("type_checking_imports", ["type_checking_imports.main"]),
-    Snapshot("type_stub", min_version=(3, 10)),
+    Snapshot("type_stubs", ["type_stubs"], min_version=(3, 10)),
     Snapshot(
         "visibility",
         render_options={

test/testdata/type_stubs.txt

@@ -0,0 +1,19 @@
+<module type_stubs  # This module has an a…
+    <function def func(x: str, y: Any, z: Iterable[str]) -> int: ...  # A simple function.>
+    <var var: list[str] = []  # Docstring override f…>
+    <class type_stubs.Class
+        <method def __init__(): ...>
+        <var attr: int = 42  # An attribute>
+        <method def meth(self, y: bool) -> bool: ...  # A simple method.>
+        <class type_stubs.Class.Subclass
+            <method def __init__(): ...>
+            <var attr: str = '42'  # An attribute>
+            <method def meth(self, y: bool) -> bool: ...  # A simple method.>
+        >
+        <method def no_type_annotation(self, z): ...  # A method not present…>
+        <method def overloaded(*args, **kwds): ...  # An overloaded method…>
+    >
+    <class type_stubs.ImportedClass  # inherited from type_stubs._utils.ImportedClass, Docstring from impor…
+        <method def __init__(): ...  # inherited from type_stubs._utils.ImportedClass.__init__>
+    >
+>

test/testdata/type_stub.py → test/testdata/type_stubs/__init__.py

@@ -2,6 +2,8 @@
 This module has an accompanying .pyi file with type stubs.
 """
 
+from ._utils import ImportedClass
+
 
 def func(x, y):
     """A simple function."""
@@ -30,3 +32,11 @@ def no_type_annotation(self, z):
 
     def overloaded(self, x):
         """An overloaded method."""
+
+
+__all__ = [
+    "func",
+    "var",
+    "Class",
+    "ImportedClass",
+]

test/testdata/type_stub.pyi → test/testdata/type_stubs/__init__.pyi

@@ -2,6 +2,8 @@ from typing import Any
 from typing import Iterable
 from typing import overload
 
+from ._utils import ImportedClass
+
 def func(x: str, y: Any, z: "Iterable[str]") -> int: ...
 
 var: list[str]
@@ -21,3 +23,10 @@ class Class:
     def overloaded(self, x: int) -> int: ...
     @overload
     def overloaded(self, x: str) -> str: ...
+
+__all__ = [
+    "func",
+    "var",
+    "Class",
+    "ImportedClass",
+]

test/testdata/type_stubs/_utils.py

@@ -0,0 +1,2 @@
+class ImportedClass:
+    """Docstring from imported py file - ideally this should be overridden."""

test/testdata/type_stubs/_utils.pyi

@@ -0,0 +1,2 @@
+class ImportedClass:
+    """Docstring from imported pyi file"""