[mypyc] fix: UnboundLocalError incorrectly raised as AttributeError

mypyc/codegen/emitfunc.py

@@ -429,19 +429,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_generated:
+                        # 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:
@@ -935,20 +937,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_generated:
+            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/lib-rt/CPy.h

@@ -853,6 +853,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

@@ -257,3 +257,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]