diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py index 286d7bfbf1079..529ccd42ef80c 100644 --- a/mypyc/irbuild/env_class.py +++ b/mypyc/irbuild/env_class.py @@ -149,16 +149,14 @@ def load_outer_env( *, borrow: bool = False, ) -> Value: - """Load the environment class for a given base into a register. + """Load the environment referenced by "base" and bind its symbols. - Additionally, iterates through all of the SymbolNode and - AssignmentTarget instances of the environment at the given index's - symtable, and adds those instances to the environment of the - current environment. This is done so that the current environment - can access outer environment variables without having to reload - all of the environment registers. + "outer_env" may contain registers and attributes from multiple environment levels. + Bind only attributes stored directly on the loaded environment, since the other targets + cannot be accessed through it. - Returns the register where the environment class was loaded. + If "borrow" is true, return a borrowed reference. In that case the environment link on + "base" must be final so that the reference remains valid. """ if borrow: assert isinstance(base.type, RInstance) @@ -167,9 +165,14 @@ def load_outer_env( assert isinstance(env.type, RInstance), f"{env} must be of type RInstance" for symbol, target in outer_env.items(): - attr_name = symbol.name - if isinstance(target, AssignmentTargetAttr): - attr_name = target.attr + # Only targets actually stored in this environment are visible through it. + if ( + not isinstance(target, AssignmentTargetAttr) + or not isinstance(target.obj.type, RInstance) + or target.obj.type.class_ir is not env.type.class_ir + ): + continue + attr_name = target.attr env.type.class_ir.attributes[attr_name] = target.type symbol_target = AssignmentTargetAttr(env, attr_name) builder.add_target(symbol, symbol_target) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 40791e8326f8d..a3e7dabf6ce20 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -26,6 +26,7 @@ BasicBlock, Branch, Call, + GetAttr, Goto, Integer, LoadErrorValue, @@ -51,7 +52,6 @@ add_vars_to_env, finalize_env_class, load_env_registers, - load_outer_env, load_outer_envs, setup_func_for_recursive_call, ) @@ -439,7 +439,8 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None: if builder.fn_info.can_merge_generator_and_env_classes(): cls.curr_env_reg = cls.self_reg else: - cls.curr_env_reg = load_outer_env(builder, cls.self_reg, builder.symtables[-1]) + cls.curr_env_reg = builder.add(GetAttr(cls.self_reg, ENV_ATTR_NAME, fitem.line)) + assert isinstance(cls.curr_env_reg.type, RInstance) # The continuation label identifies where execution resumes when the generator is next # advanced. Only the serialized generator helper accesses it, so keep it on the private diff --git a/mypyc/test/test_generator_spills.py b/mypyc/test/test_generator_spills.py index 440ebb068675e..9e3a0fb76e164 100644 --- a/mypyc/test/test_generator_spills.py +++ b/mypyc/test/test_generator_spills.py @@ -7,9 +7,16 @@ import unittest -from mypyc.common import GENERATOR_ATTRIBUTE_PREFIX, TEMP_ATTR_NAME +from mypyc.common import ( + ENV_ATTR_NAME, + GENERATOR_ATTRIBUTE_PREFIX, + NEXT_LABEL_ATTR_NAME, + SELF_NAME, + TEMP_ATTR_NAME, +) from mypyc.ir.class_ir import ClassIR from mypyc.ir.ops import Assign, GetAttr +from mypyc.ir.rtypes import RInstance from mypyc.test.testutil import build_ir_for_single_file2 from mypyc.transform.generator_spills import registers_live_across_yield @@ -39,6 +46,29 @@ def frame_variables(cl: ClassIR) -> set[str]: class TestGeneratorSpills(unittest.TestCase): + def test_generator_environment_excludes_helper_arguments(self) -> None: + # Loading the outer environment must not add the generator helper's arguments or + # state to the environment class, or overwrite the type of its self-reference. + module, _, _, _ = build_ir_for_single_file2("""\ +from typing import Iterator + +def outer() -> Iterator[str]: + captured = "value" + def nested() -> str: + return captured + yield nested() +""".splitlines()) + frame = next(cl for cl in module.classes if cl.has_running_flag) + env_type = frame.attributes[ENV_ATTR_NAME] + assert isinstance(env_type, RInstance) + environment = env_type.class_ir + + self_type = environment.attributes[SELF_NAME] + assert isinstance(self_type, RInstance) + assert self_type.class_ir is environment + for name in ("type", "value", "traceback", "arg", "stop_iter_ptr", NEXT_LABEL_ATTR_NAME): + assert name not in environment.attributes + def test_saved_exception_state_is_promoted(self) -> None: cl = generator_class( """\