recursively evaluate TYPE_CHECKING blocks, fix #648 (#649)

CHANGELOG.md

@@ -6,6 +6,9 @@
 
  - pdoc now documents PyO3 or pybind11 submodules that are not picked up by Python's builtin pkgutil module.
    ([#633](https://github.com/mitmproxy/pdoc/issues/633), @mhils)
+ - Imports in a TYPE_CHECKING section that reference members defined in another module's TYPE_CHECKING section now work
+   correctly.
+   ([#649](https://github.com/mitmproxy/pdoc/pull/649), @mhils)
  - pdoc now supports Python 3.12's `type` statements and has improved `TypeAlias` rendering.
    ([#651](https://github.com/mitmproxy/pdoc/pull/651), @mhils)
  - Add support for `code-block` ReST directives

pdoc/doc_types.py

@@ -124,9 +124,9 @@ def safe_eval_type(
 
     # Simple _eval_type has failed. We now execute all TYPE_CHECKING sections in the module and try again.
     if module:
+        assert module.__dict__ is globalns
         try:
-            code = compile(type_checking_sections(module), "<string>", "exec")
-            eval(code, globalns, globalns)
+            _eval_type_checking_sections(module, set())
         except Exception as e:
             warnings.warn(
                 f"Failed to run TYPE_CHECKING code while parsing {t} type annotation for {fullname}: {e}"
@@ -148,7 +148,34 @@ def safe_eval_type(
             f"Error parsing type annotation {t} for {fullname}. Import of {mod} failed: {err}"
         )
         return t
-    return safe_eval_type(t, {mod: val, **globalns}, localns, module, fullname)
+    else:
+        globalns[mod] = val
+    return safe_eval_type(t, globalns, localns, module, fullname)
+
+
+def _eval_type_checking_sections(module: types.ModuleType, seen: set) -> None:
+    """
+    Evaluate all TYPE_CHECKING sections within a module.
+
+    The added complication here is that TYPE_CHECKING sections may import members from other modules' TYPE_CHECKING
+    sections. So we try to recursively execute those other modules' TYPE_CHECKING sections as well.
+    See https://github.com/mitmproxy/pdoc/issues/648 for a real world example.
+    """
+    if module.__name__ in seen:
+        raise RecursionError(f"Recursion error when importing {module.__name__}.")
+    seen.add(module.__name__)
+
+    code = compile(type_checking_sections(module), "<string>", "exec")
+    while True:
+        try:
+            eval(code, module.__dict__, module.__dict__)
+        except ImportError as e:
+            if e.name is not None and (mod := sys.modules.get(e.name, None)):
+                _eval_type_checking_sections(mod, seen)
+            else:
+                raise
+        else:
+            break
 
 
 def _eval_type(t, globalns, localns, recursive_guard=frozenset()):

test/test_doc_types.py

@@ -12,8 +12,9 @@
     "typestr", ["totally_unknown_module", "!!!!", "html.unknown_attr"]
 )
 def test_eval_fail(typestr):
+    a = types.ModuleType("a")
     with pytest.warns(UserWarning, match="Error parsing type annotation"):
-        assert safe_eval_type(typestr, {}, None, types.ModuleType("a"), "a") == typestr
+        assert safe_eval_type(typestr, a.__dict__, None, a, "a") == typestr
 
 
 def test_eval_fail2(monkeypatch):
@@ -22,8 +23,9 @@ def test_eval_fail2(monkeypatch):
         "get_source",
         lambda _: "import typing\nif typing.TYPE_CHECKING:\n\traise RuntimeError()",
     )
+    a = types.ModuleType("a")
     with pytest.warns(UserWarning, match="Failed to run TYPE_CHECKING code"):
-        assert safe_eval_type("xyz", {}, None, types.ModuleType("a"), "a") == "xyz"
+        assert safe_eval_type("xyz", a.__dict__, None, a, "a") == "xyz"
 
 
 def test_eval_fail3(monkeypatch):
@@ -32,16 +34,24 @@ def test_eval_fail3(monkeypatch):
         "get_source",
         lambda _: "import typing\nif typing.TYPE_CHECKING:\n\tFooFn = typing.Callable[[],int]",
     )
+    a = types.ModuleType("a")
+    a.__dict__["typing"] = typing
     with pytest.warns(
         UserWarning,
         match="Error parsing type annotation .+ after evaluating TYPE_CHECKING blocks",
     ):
-        assert (
-            safe_eval_type(
-                "FooFn[int]", {"typing": typing}, None, types.ModuleType("a"), "a"
-            )
-            == "FooFn[int]"
-        )
+        assert safe_eval_type("FooFn[int]", a.__dict__, None, a, "a") == "FooFn[int]"
+
+
+def test_eval_fail_import_nonexistent(monkeypatch):
+    monkeypatch.setattr(
+        doc_ast,
+        "get_source",
+        lambda _: "import typing\nif typing.TYPE_CHECKING:\n\timport nonexistent_module",
+    )
+    a = types.ModuleType("a")
+    with pytest.warns(UserWarning, match="No module named 'nonexistent_module'"):
+        assert safe_eval_type("xyz", a.__dict__, None, a, "a") == "xyz"
 
 
 def test_eval_union_types_on_old_python(monkeypatch):
@@ -53,3 +63,21 @@ def test_eval_union_types_on_old_python(monkeypatch):
     ):
         # str never implements `|`, so we can use that to trigger the error on newer versions.
         safe_eval_type('"foo" | "bar"', {}, None, None, "example")
+
+
+def test_recurse(monkeypatch):
+    def get_source(mod):
+        if mod == a:
+            return "import typing\nif typing.TYPE_CHECKING:\n\tfrom b import Foo"
+        else:
+            return "import typing\nif typing.TYPE_CHECKING:\n\tfrom a import Foo"
+
+    a = types.ModuleType("a")
+    b = types.ModuleType("b")
+
+    monkeypatch.setattr(doc_ast, "get_source", get_source)
+    monkeypatch.setitem(sys.modules, "a", a)
+    monkeypatch.setitem(sys.modules, "b", b)
+
+    with pytest.warns(UserWarning, match="Recursion error when importing a"):
+        assert safe_eval_type("xyz", a.__dict__, None, a, "a") == "xyz"

test/test_snapshot.py

@@ -162,7 +162,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"),
+    Snapshot("type_checking_imports", ["type_checking_imports.main"]),
     Snapshot("type_stub", min_version=(3, 10)),
     Snapshot(
         "visibility",