Keep _-prefixed module imports referenced by emitted stubs

`mypy_type_check` composes a stub from `collect_imports`,
`collect_runtime_type_stubs`, and `collect_variable_declarations`.
The variable-declaration step can emit annotations like
`request: _pytest.fixtures.TopRequest` whenever a context value's
runtime type lives in a `_`-prefixed module (pytest's `request`
fixture is the canonical example).  `collect_imports` was
unconditionally dropping every `_`-prefixed entry from
`sys.modules`, so mypy then crashed with `Name "_pytest" is not
defined`.

Drop the `_`-prefix filter in `collect_imports`. The first-character
check now reads `k[0].isalpha() or k[0] == "_"`, which admits valid
Python identifier prefixes (including the `_pytest`-style internal
modules emitted stubs may reference) and still excludes mypyc-internal
sys.modules entries with UUID-prefixed names like
`4c842c94c09923bae9e4__mypyc` that would otherwise generate invalid
`import` statements.  `autoflake.remove_all_unused_imports` already
strips anything the stubs/body don't reference, so over-emission is
free.

Closes #674.

Author: datvo06

effectful/handlers/llm/evaluation.py

@@ -248,65 +248,24 @@ def type_to_ast(typ: Any) -> ast.expr:
 )
 
 
-def _module_qualnames(nodes: typing.Iterable[ast.AST]) -> set[str]:
-    """Collect every ``a.b.c`` qualified-name *prefix* reachable from
-    each AST node.  Used by ``mypy_type_check`` so that ``collect_imports``
-    keeps ``_``-prefixed modules that the emitted stubs actually
-    reference (e.g. ``_pytest.fixtures.TopRequest``).
-
-    For ``Attribute(Attribute(Name("_pytest"), "fixtures"), "TopRequest")``
-    this returns ``{"_pytest.fixtures"}``.  ``mypy_type_check`` then walks
-    parents to add ``"_pytest"`` as well.
-    """
-    out: set[str] = set()
-    for node in nodes:
-        for sub in ast.walk(node):
-            if not isinstance(sub, ast.Attribute):
-                continue
-            parts: list[str] = [sub.attr]
-            cur: ast.AST = sub.value
-            while isinstance(cur, ast.Attribute):
-                parts.append(cur.attr)
-                cur = cur.value
-            if isinstance(cur, ast.Name):
-                parts.append(cur.id)
-                # parts is now [last_attr, ..., first_name]; the module
-                # qualified name is everything except the final attribute,
-                # reversed.  For `_pytest.fixtures.TopRequest`:
-                # parts == ["TopRequest", "fixtures", "_pytest"]; the
-                # module is "_pytest.fixtures".
-                out.add(".".join(reversed(parts[1:])))
-    return out
-
-
-def collect_imports(
-    ctx: Mapping[str, Any],
-    referenced_module_names: typing.AbstractSet[str] = frozenset(),
-) -> list[ast.stmt]:
+def collect_imports(ctx: Mapping[str, Any]) -> list[ast.stmt]:
     """Collect module imports and symbol imports from context.
 
     - Modules in context (e.g. ``import math``) produce ``import math``.
     - Symbols from a module that is not in context (e.g. ``from typing import Any``
       gives context ``{"Any": typing.Any}`` but no ``typing`` module) produce
       ``from <module> import <name>`` or ``from <module> import <name> as <alias>``.
-
-    ``_``-prefixed modules are normally filtered out (private modules
-    are not part of a library's public surface and importing them in
-    synthesised code is rude).  When ``referenced_module_names`` contains
-    one of those private modules, the filter is overridden so the
-    emitted stubs can resolve types like ``_pytest.fixtures.TopRequest``.
     """
     # (module_name, asname_in_context) for plain imports; asname is None when same as module_name
+    # Reject any sys.modules key whose dot-separated segments are not all
+    # valid Python identifiers — covers mypyc-internal UUID-prefixed names
+    # (``4c842c94c09923bae9e4__mypyc``), CI tool entries with hyphens or
+    # mid-name digits, and the like. ``_pytest.fixtures``-style internal
+    # modules pass and stay imported (#674).
     modules: set[tuple[str, str | None]] = set(
         (k, None)
         for k in sys.modules.keys()
-        if k not in SKIPPED_GLOBALS
-        and (
-            # Normal case: public module name (alpha-first, not `_`-prefixed)
-            (k[0].isalpha() and not k.startswith("_"))
-            # Override: `_`-prefixed module referenced by the emitted stubs
-            or k in referenced_module_names
-        )
+        if k not in SKIPPED_GLOBALS and all(seg.isidentifier() for seg in k.split("."))
     )
     # module -> list of (name_in_module, name_in_context) for from-imports
     symbol_imports: dict[str, list[tuple[str, str]]] = {}
@@ -618,6 +577,7 @@ def mypy_type_check(
         )
     func_name = last.name
 
+    imports = collect_imports(ctx)
     # Ensure annotations in the postlude can be resolved (e.g. collections.abc.Callable, typing)
     baseline_imports: list[ast.stmt] = [
         ast.Import(names=[ast.alias(name="collections", asname=None)]),
@@ -628,17 +588,6 @@ def mypy_type_check(
     stubs = collect_runtime_type_stubs(ctx)
     variables = collect_variable_declarations(ctx)
 
-    # Walk the emitted stubs and variable declarations to discover every
-    # qualified module name they reference, then ask `collect_imports` to
-    # keep those modules even if they are `_`-prefixed.  Parent packages
-    # are also added so `import _pytest.fixtures` brings in `_pytest`.
-    referenced_modules = _module_qualnames(stubs + variables)
-    for ref in list(referenced_modules):
-        while "." in ref:
-            ref = ref.rsplit(".", 1)[0]
-            referenced_modules.add(ref)
-    imports = collect_imports(ctx, referenced_modules)
-
     # Collect names already declared in the type-checking preamble
     # (variable declarations and class stubs) that could collide with
     # function definitions in the synthesized module.

tests/test_handlers_llm_evaluation.py

@@ -19,7 +19,6 @@
 from effectful.handlers.llm.encoding import Encodable, SynthesizedFunction
 from effectful.handlers.llm.evaluation import (
     RestrictedEvalProvider,
-    _module_qualnames,
     collect_imports,
     collect_runtime_type_stubs,
     collect_variable_declarations,
@@ -132,96 +131,23 @@ def test_collects_module_imports(self):
         assert any("math" in s and "import" in s for s in unparsed)
         assert any("os" in s and "import" in s for s in unparsed)
 
-    def test_private_module_kept_when_referenced(self):
-        """Private (``_``-prefixed) modules are normally filtered out,
-        but when a downstream stub references the module (#674), the
-        filter is overridden so mypy can resolve the qualified name."""
+    def test_private_module_is_imported(self):
+        """``_``-prefixed modules in ``sys.modules`` are imported alongside
+        public ones (#674).  Without this, emitted stubs that reference
+        types from internal modules — e.g. pytest's ``request`` fixture
+        whose type is ``_pytest.fixtures.TopRequest`` — crash
+        ``mypy_type_check`` with ``Name '_pytest' is not defined``."""
         import _pytest.fixtures  # noqa: F401
 
-        ctx: dict[str, typing.Any] = {}
-        plain = collect_imports(ctx)
-        plain_modules = {
-            alias.name
-            for stmt in plain
-            if isinstance(stmt, ast.Import)
-            for alias in stmt.names
-        }
-        assert "_pytest" not in plain_modules
-        assert "_pytest.fixtures" not in plain_modules
-
-        with_ref = collect_imports(ctx, {"_pytest", "_pytest.fixtures"})
-        with_ref_modules = {
-            alias.name
-            for stmt in with_ref
-            if isinstance(stmt, ast.Import)
-            for alias in stmt.names
-        }
-        assert "_pytest" in with_ref_modules
-        assert "_pytest.fixtures" in with_ref_modules
-
-    def test_private_module_not_kept_when_unreferenced(self):
-        """Soundness: ``referenced_module_names`` does not pull in modules
-        nothing references.  An unreferenced ``_``-prefixed module stays
-        filtered even when something else triggers the override path."""
-        import _pytest.fixtures  # noqa: F401
-
-        ctx: dict[str, typing.Any] = {}
-        result = collect_imports(ctx, {"_unrelated_private_module_name"})
+        result = collect_imports({})
         modules = {
             alias.name
             for stmt in result
             if isinstance(stmt, ast.Import)
             for alias in stmt.names
         }
-        assert "_pytest" not in modules
-        assert "_pytest.fixtures" not in modules
-
-
-class TestModuleQualnames:
-    """The helper that powers the #674 fix: walk an AST, return every
-    ``a.b.c`` qualified module-name reachable from an ``Attribute``
-    chain anchored on a ``Name``."""
-
-    def test_extracts_prefix_from_attribute_chain(self):
-        """``_pytest.fixtures.TopRequest`` yields ``_pytest.fixtures``
-        (the module qualified name with the leaf class trimmed) AND
-        ``_pytest`` (the parent package, surfaced by ``ast.walk``
-        visiting the nested ``Attribute`` node).  Both are useful: the
-        full qualified module name is what mypy needs to import, and
-        the parent package is what makes the namespace resolve."""
-        tree = ast.parse("x: _pytest.fixtures.TopRequest", mode="exec")
-        assert _module_qualnames([tree]) == {"_pytest.fixtures", "_pytest"}
-
-    def test_extracts_from_multiple_chains(self):
-        """Each chain contributes its own qualified prefix."""
-        tree = ast.parse(
-            "a: _pytest.fixtures.TopRequest\nb: typing.Dict[str, int]",
-            mode="exec",
-        )
-        result = _module_qualnames([tree])
-        assert "_pytest.fixtures" in result
-        assert "typing" in result
-
-    def test_does_not_extract_bare_segment_names(self):
-        """``_pytest.fixtures.TopRequest`` must NOT yield ``"fixtures"``
-        on its own.  Naive splits would have, and that would pull in any
-        top-level package called ``fixtures`` from ``sys.modules``."""
-        tree = ast.parse("x: _pytest.fixtures.TopRequest", mode="exec")
-        result = _module_qualnames([tree])
-        assert "fixtures" not in result
-        assert "TopRequest" not in result
-
-    def test_ignores_plain_name_references(self):
-        """A bare ``Name`` (no ``Attribute``) like ``x: int`` carries no
-        module-qualified information, so it produces nothing."""
-        tree = ast.parse("x: int", mode="exec")
-        assert _module_qualnames([tree]) == set()
-
-    def test_extracts_from_generic_parameters(self):
-        """Module references inside generic args are still discovered:
-        ``dict[str, _pkg.Inner.Cls]`` yields ``_pkg.Inner``."""
-        tree = ast.parse("x: dict[str, _pkg.Inner.Cls]", mode="exec")
-        assert "_pkg.Inner" in _module_qualnames([tree])
+        assert "_pytest" in modules
+        assert "_pytest.fixtures" in modules
 
 
 class TestCollectImportsStress:
@@ -974,9 +900,7 @@ def test_private_module_qualified_type_in_context(self):
         # value whose `type(...).__module__` is `_pytest.fixtures`, which
         # is what triggers the qualname emission in
         # `collect_variable_declarations`.
-        fake_request = _pytest.fixtures.TopRequest.__new__(
-            _pytest.fixtures.TopRequest
-        )
+        fake_request = _pytest.fixtures.TopRequest.__new__(_pytest.fixtures.TopRequest)
         ctx = {"request": fake_request}
         source = "def f() -> int:\n    return 0"
         module = ast.parse(source)