Keep _-prefixed module imports referenced by emitted stubs (#674)

effectful/handlers/llm/evaluation.py

@@ -257,10 +257,15 @@ def collect_imports(ctx: Mapping[str, Any]) -> list[ast.stmt]:
       ``from <module> import <name>`` or ``from <module> import <name> as <alias>``.
     """
     # (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 not k.startswith("_") and k[0].isalpha()
+        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]]] = {}

tests/test_handlers_llm_evaluation.py

@@ -131,6 +131,24 @@ 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_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
+
+        result = collect_imports({})
+        modules = {
+            alias.name
+            for stmt in result
+            if isinstance(stmt, ast.Import)
+            for alias in stmt.names
+        }
+        assert "_pytest" in modules
+        assert "_pytest.fixtures" in modules
+
 
 class TestCollectImportsStress:
     """Stress test collect_imports with get_context: imports, aliases, external symbols."""
@@ -870,6 +888,25 @@ def test_simple_function_with_get_context(self):
         module = ast.parse(source)
         mypy_type_check(module, get_context(), [int, str], bool)
 
+    def test_private_module_qualified_type_in_context(self):
+        """Regression for #674: a ctx value whose runtime type lives in
+        a ``_``-prefixed module must not crash ``mypy_type_check`` with
+        ``Name '_pytest' is not defined``.  Uses a pytest fixture-request
+        instance because that is the path that surfaced the bug."""
+        import _pytest.fixtures
+
+        # Build a `request`-shaped instance.  We cannot easily instantiate
+        # `TopRequest` properly outside pytest, but `__new__` gives us a
+        # 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)
+        ctx = {"request": fake_request}
+        source = "def f() -> int:\n    return 0"
+        module = ast.parse(source)
+        # Must not raise.
+        mypy_type_check(module, ctx, None, int)
+
     def test_simple_function_no_params_with_get_context(self):
         """Function with no params, returns int; get_context()."""
         _ = 1  # noqa: F841