fix UnboundLocalError handling for env classes (#9)

Author: BobTheBuidler

mypyc/codegen/emitfunc.py

@@ -419,7 +419,7 @@ def visit_get_attr(self, op: GetAttr) -> None:
                         self.emitter.emit_line("}")
                 if not merged_branch:
                     var_name = op.attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX)
-                    if cl.is_generated:
+                    if cl.is_environment:
                         # A generated class does not "exist" to the user, this is just an unbound
                         # variable in their code, not a missing attribute on the generated class.
                         exc_class = "PyExc_UnboundLocalError"
@@ -924,7 +924,7 @@ def emit_traceback(self, op: Branch) -> None:
     def emit_attribute_error(self, op: Branch, class_ir: ClassIR, attr: str) -> None:
         assert op.traceback_entry is not None
         globals_static = self.emitter.static_name("globals", self.module_name)
-        if class_ir.is_generated:
+        if class_ir.is_environment:
             self.emit_line(
                 'CPy_UnboundLocalError("%s", "%s", "%s", %d, %s);'
                 % (

mypyc/ir/class_ir.py

@@ -91,6 +91,7 @@ def __init__(
         module_name: str,
         is_trait: bool = False,
         is_generated: bool = False,
+        is_environment: bool = False,
         is_abstract: bool = False,
         is_ext_class: bool = True,
         is_final_class: bool = False,
@@ -99,6 +100,8 @@ def __init__(
         self.module_name = module_name
         self.is_trait = is_trait
         self.is_generated = is_generated
+        # True for env classes where attributes represent locals.
+        self.is_environment = is_environment
         self.is_abstract = is_abstract
         self.is_ext_class = is_ext_class
         self.is_final_class = is_final_class
@@ -231,6 +234,7 @@ def __repr__(self) -> str:
             "ClassIR("
             "name={self.name}, module_name={self.module_name}, "
             "is_trait={self.is_trait}, is_generated={self.is_generated}, "
+            "is_environment={self.is_environment}, "
             "is_abstract={self.is_abstract}, is_ext_class={self.is_ext_class}, "
             "is_final_class={self.is_final_class}"
             ")".format(self=self)
@@ -380,6 +384,7 @@ def serialize(self) -> JsonDict:
             "is_ext_class": self.is_ext_class,
             "is_abstract": self.is_abstract,
             "is_generated": self.is_generated,
+            "is_environment": self.is_environment,
             "is_augmented": self.is_augmented,
             "is_final_class": self.is_final_class,
             "inherits_python": self.inherits_python,
@@ -438,6 +443,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
 
         ir.is_trait = data["is_trait"]
         ir.is_generated = data["is_generated"]
+        ir.is_environment = data.get("is_environment", False)
         ir.is_abstract = data["is_abstract"]
         ir.is_ext_class = data["is_ext_class"]
         ir.is_augmented = data["is_augmented"]

mypyc/irbuild/env_class.py

@@ -54,6 +54,7 @@ class is generated, the function environment has not yet been
         is_generated=True,
         is_final_class=True,
     )
+    env_class.is_environment = True
     env_class.reuse_freed_instance = True
     env_class.attributes[SELF_NAME] = RInstance(env_class)
     if builder.fn_info.is_nested:

mypyc/irbuild/generator.py

@@ -164,6 +164,7 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR:
     generator_class_ir = mapper.fdef_to_generator[builder.fn_info.fitem]
     if builder.fn_info.can_merge_generator_and_env_classes():
         builder.fn_info.env_class = generator_class_ir
+        generator_class_ir.is_environment = True
     else:
         generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class)