Add prefix also to local vars in generator functions

Author: p-sawicki

mypyc/codegen/emitfunc.py

@@ -8,6 +8,7 @@
 from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.emit import DEBUG_ERRORS, Emitter, TracebackAndGotoHandler, c_array_initializer
 from mypyc.common import (
+    GENERATOR_ATTRIBUTE_PREFIX,
     HAVE_IMMORTAL,
     MODULE_PREFIX,
     NATIVE_PREFIX,
@@ -436,7 +437,9 @@ def visit_get_attr(self, op: GetAttr) -> None:
                     exc_class = "PyExc_AttributeError"
                     self.emitter.emit_line(
                         'PyErr_SetString({}, "attribute {} of {} undefined");'.format(
-                            exc_class, repr(op.attr), repr(cl.name)
+                            exc_class,
+                            repr(op.attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX)),
+                            repr(cl.name),
                         )
                     )
 
@@ -938,7 +941,7 @@ def emit_attribute_error(self, op: Branch, class_name: str, attr: str) -> None:
                 self.source_path.replace("\\", "\\\\"),
                 op.traceback_entry[0],
                 class_name,
-                attr,
+                attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX),
                 op.traceback_entry[1],
                 globals_static,
             )

mypyc/common.py

@@ -22,6 +22,7 @@
 LAMBDA_NAME: Final = "__mypyc_lambda__"
 PROPSET_PREFIX: Final = "__mypyc_setter__"
 SELF_NAME: Final = "__mypyc_self__"
+GENERATOR_ATTRIBUTE_PREFIX: Final = "__mypyc_generator_attribute__"
 
 # Max short int we accept as a literal is based on 32-bit platforms,
 # so that we can just always emit the same code.

mypyc/irbuild/builder.py

@@ -59,7 +59,7 @@
 )
 from mypy.util import module_prefix, split_target
 from mypy.visitor import ExpressionVisitor, StatementVisitor
-from mypyc.common import BITMAP_BITS, SELF_NAME, TEMP_ATTR_NAME
+from mypyc.common import BITMAP_BITS, GENERATOR_ATTRIBUTE_PREFIX, SELF_NAME, TEMP_ATTR_NAME
 from mypyc.crash import catch_errors
 from mypyc.errors import Errors
 from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
@@ -651,7 +651,11 @@ def get_assignment_target(
                     # current environment.
                     if self.fn_info.is_generator:
                         return self.add_var_to_env_class(
-                            symbol, reg_type, self.fn_info.generator_class, reassign=False
+                            symbol,
+                            reg_type,
+                            self.fn_info.generator_class,
+                            reassign=False,
+                            prefix=GENERATOR_ATTRIBUTE_PREFIX,
                         )
 
                     # Otherwise define a new local variable.

mypyc/irbuild/env_class.py

@@ -235,7 +235,7 @@ def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None:
                 # will generate different callable classes, so the callable
                 # class that gets instantiated must be generic.
                 builder.add_var_to_env_class(
-                    nested_fn, object_rprimitive, env_for_func, reassign=False
+                    nested_fn, object_rprimitive, env_for_func, reassign=False, prefix=prefix
                 )
 
 

mypyc/irbuild/generator.py

@@ -13,7 +13,7 @@
 from typing import Callable
 
 from mypy.nodes import ARG_OPT, FuncDef, Var
-from mypyc.common import ENV_ATTR_NAME, NEXT_LABEL_ATTR_NAME
+from mypyc.common import ENV_ATTR_NAME, GENERATOR_ATTRIBUTE_PREFIX, NEXT_LABEL_ATTR_NAME
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FuncDecl, FuncIR
 from mypyc.ir.ops import (
@@ -59,8 +59,6 @@
     restore_exc_info_op,
 )
 
-GENERATOR_ATTRIBUTE_PREFIX = "__mypyc_generator_attribute__"
-
 
 def gen_generator_func(
     builder: IRBuilder,

mypyc/test-data/run-generators.test

@@ -871,3 +871,30 @@ def test_undefined_int_in_environment() -> None:
 
     with assertRaises(AttributeError):  # TODO: Should be UnboundLocalError
       list(gen2(False))
+
+[case testVariableWithSameNameAsHelperMethod]
+from testutil import assertRaises
+from typing import Iterator
+
+def gen_send() -> Iterator[int]:
+    send = 1
+    yield send + 1
+
+def gen_throw() -> Iterator[int]:
+    throw = 42
+    yield throw * 2
+
+def undefined() -> Iterator[int]:
+    if int():
+        send = 1
+    yield send + 1
+
+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.
+        list(undefined())