[mypyc] fix: UnboundLocalError incorrectly raised as AttributeError

mypyc/codegen/emitfunc.py

@@ -413,19 +413,21 @@ def visit_get_attr(self, op: GetAttr) -> None:
                     ):
                         # Generate code for the following branch here to avoid
                         # redundant branches in the generated code.
-                        self.emit_attribute_error(branch, cl.name, op.attr)
+                        self.emit_attribute_error(branch, cl, op.attr)
                         self.emit_line("goto %s;" % self.label(branch.true))
                         merged_branch = branch
                         self.emitter.emit_line("}")
                 if not merged_branch:
-                    exc_class = "PyExc_AttributeError"
-                    self.emitter.emit_line(
-                        'PyErr_SetString({}, "attribute {} of {} undefined");'.format(
-                            exc_class,
-                            repr(op.attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX)),
-                            repr(cl.name),
-                        )
-                    )
+                    var_name = op.attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX)
+                    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"
+                        exc_msg = f"local variable {var_name!r} referenced before assignment"
+                    else:
+                        exc_class = "PyExc_AttributeError"
+                        exc_msg = f"attribute {var_name!r} of {cl.name!r} undefined"
+                    self.emitter.emit_line(f'PyErr_SetString({exc_class}, "{exc_msg}");')
 
             if attr_rtype.is_refcounted and not op.is_borrowed:
                 if not merged_branch and not always_defined:
@@ -919,20 +921,32 @@ def emit_traceback(self, op: Branch) -> None:
         if op.traceback_entry is not None:
             self.emitter.emit_traceback(self.source_path, self.module_name, op.traceback_entry)
 
-    def emit_attribute_error(self, op: Branch, class_name: str, attr: str) -> 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)
-        self.emit_line(
-            'CPy_AttributeError("%s", "%s", "%s", "%s", %d, %s);'
-            % (
-                self.source_path.replace("\\", "\\\\"),
-                op.traceback_entry[0],
-                class_name,
-                attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX),
-                op.traceback_entry[1],
-                globals_static,
+        if class_ir.is_environment:
+            self.emit_line(
+                'CPy_UnboundLocalError("%s", "%s", "%s", %d, %s);'
+                % (
+                    self.source_path.replace("\\", "\\\\"),
+                    op.traceback_entry[0],
+                    attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX),
+                    op.traceback_entry[1],
+                    globals_static,
+                )
+            )
+        else:
+            self.emit_line(
+                'CPy_AttributeError("%s", "%s", "%s", "%s", %d, %s);'
+                % (
+                    self.source_path.replace("\\", "\\\\"),
+                    op.traceback_entry[0],
+                    class_ir.name,
+                    attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX),
+                    op.traceback_entry[1],
+                    globals_static,
+                )
             )
-        )
         if DEBUG_ERRORS:
             self.emit_line('assert(PyErr_Occurred() != NULL && "failure w/o err!");')
 

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)
 

mypyc/lib-rt/CPy.h

@@ -861,6 +861,8 @@ void CPy_TypeErrorTraceback(const char *filename, const char *funcname, int line
                             PyObject *globals, const char *expected, PyObject *value);
 void CPy_AttributeError(const char *filename, const char *funcname, const char *classname,
                         const char *attrname, int line, PyObject *globals);
+void CPy_UnboundLocalError(const char *filename, const char *funcname, const char *attrname,
+                           int line, PyObject *globals);
 
 
 // Misc operations

mypyc/lib-rt/exc_ops.c

@@ -259,3 +259,11 @@ void CPy_AttributeError(const char *filename, const char *funcname, const char *
     PyErr_SetString(PyExc_AttributeError, buf);
     CPy_AddTraceback(filename, funcname, line, globals);
 }
+
+void CPy_UnboundLocalError(const char *filename, const char *funcname, const char *attrname,
+                           int line, PyObject *globals) {
+    char buf[500];
+    snprintf(buf, sizeof(buf), "local variable '%.200s' referenced before assignment", attrname);
+    PyErr_SetString(PyExc_UnboundLocalError, buf);
+    CPy_AddTraceback(filename, funcname, line, globals);
+}

mypyc/test-data/run-generators.test

@@ -866,7 +866,7 @@ def test_bitmap_is_cleared_when_object_is_reused() -> None:
     list(gen(True))
 
     # Ensure bitmap has been cleared.
-    with assertRaises(AttributeError):  # TODO: Should be UnboundLocalError
+    with assertRaises(UnboundLocalError):
       list(gen(False))
 
 def gen2(set: bool) -> Iterator[int]:
@@ -878,7 +878,7 @@ def gen2(set: bool) -> Iterator[int]:
 def test_undefined_int_in_environment() -> None:
     list(gen2(True))
 
-    with assertRaises(AttributeError):  # TODO: Should be UnboundLocalError
+    with assertRaises(UnboundLocalError):
       list(gen2(False))
 
 [case testVariableWithSameNameAsHelperMethod]
@@ -902,10 +902,9 @@ def test_same_names() -> None:
     assert list(gen_send()) == [2]
     assert list(gen_throw()) == [84]
 
-    with assertRaises(AttributeError, "attribute 'send' of 'undefined_gen' undefined"):
-        # TODO: Should be UnboundLocalError, this test verifies that the attribute name
-        # matches the variable name in the input code, since internally it's generated
-        # with a prefix.
+    with assertRaises(UnboundLocalError, "local variable 'send' referenced before assignment"):
+        # this test verifies that the attribute name matches the variable name
+        # in the input code, since internally it's generated with a prefix.
         list(undefined())
 
 [case testGeneratorInheritance]