From e43739b92a534b75c0acf313c8467815ef00577c Mon Sep 17 00:00:00 2001 From: datvo06 Date: Sun, 7 Jun 2026 22:19:13 -0400 Subject: [PATCH 1/2] 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`. Fix shape (#674): - New helper `_module_qualnames(nodes)` walks an AST and returns every `a.b.c` module-qualified prefix anchored on an `Attribute` chain. Bare segment names (like "fixtures" from `_pytest.fixtures.TopRequest`) are explicitly NOT yielded, so the fix can't accidentally pull in an unrelated top-level package. - `collect_imports` takes an optional `referenced_module_names` set. Modules from `sys.modules` are kept either if they are normal public names OR they appear in the referenced set. - `mypy_type_check` computes `stubs` and `variables` first, walks them with `_module_qualnames`, adds parent packages, then calls `collect_imports` with the discovered names. Closes #674. --- effectful/handlers/llm/evaluation.py | 62 +++++++++++++- tests/test_handlers_llm_evaluation.py | 113 ++++++++++++++++++++++++++ 2 files changed, 172 insertions(+), 3 deletions(-) diff --git a/effectful/handlers/llm/evaluation.py b/effectful/handlers/llm/evaluation.py index b4c4ecf67..cb1595ef2 100644 --- a/effectful/handlers/llm/evaluation.py +++ b/effectful/handlers/llm/evaluation.py @@ -248,19 +248,65 @@ def type_to_ast(typ: Any) -> ast.expr: ) -def collect_imports(ctx: Mapping[str, Any]) -> list[ast.stmt]: +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]: """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 import `` or ``from import as ``. + + ``_``-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 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 ( + # 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 + ) ) # module -> list of (name_in_module, name_in_context) for from-imports symbol_imports: dict[str, list[tuple[str, str]]] = {} @@ -572,7 +618,6 @@ 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)]), @@ -583,6 +628,17 @@ 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. diff --git a/tests/test_handlers_llm_evaluation.py b/tests/test_handlers_llm_evaluation.py index bf064732f..9a9706193 100644 --- a/tests/test_handlers_llm_evaluation.py +++ b/tests/test_handlers_llm_evaluation.py @@ -19,6 +19,7 @@ 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, @@ -131,6 +132,97 @@ 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.""" + 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"}) + 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]) + class TestCollectImportsStress: """Stress test collect_imports with get_context: imports, aliases, external symbols.""" @@ -870,6 +962,27 @@ 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 From d4143d2360e72089bf04ef38b21dede05c968fdd Mon Sep 17 00:00:00 2001 From: datvo06 Date: Sun, 7 Jun 2026 22:23:48 -0400 Subject: [PATCH 2/2] 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. --- effectful/handlers/llm/evaluation.py | 67 +++---------------- tests/test_handlers_llm_evaluation.py | 96 +++------------------------ 2 files changed, 18 insertions(+), 145 deletions(-) diff --git a/effectful/handlers/llm/evaluation.py b/effectful/handlers/llm/evaluation.py index cb1595ef2..224f79e5d 100644 --- a/effectful/handlers/llm/evaluation.py +++ b/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 import `` or ``from import as ``. - - ``_``-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. diff --git a/tests/test_handlers_llm_evaluation.py b/tests/test_handlers_llm_evaluation.py index 9a9706193..91b58ec14 100644 --- a/tests/test_handlers_llm_evaluation.py +++ b/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)