[mypyc] Add prefix to attributes of generator classes

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.
@@ -1333,10 +1337,11 @@ def add_var_to_env_class(
         base: FuncInfo | ImplicitClass,
         reassign: bool = False,
         always_defined: bool = False,
+        prefix: str = "",
     ) -> AssignmentTarget:
         # First, define the variable name as an attribute of the environment class, and then
         # construct a target for that attribute.
-        name = remangle_redefinition_name(var.name)
+        name = prefix + remangle_redefinition_name(var.name)
         self.fn_info.env_class.attributes[name] = rtype
         if always_defined:
             self.fn_info.env_class.attrs_with_defaults.add(name)

mypyc/irbuild/env_class.py

@@ -18,7 +18,13 @@ def g() -> int:
 from __future__ import annotations
 
 from mypy.nodes import Argument, FuncDef, SymbolNode, Var
-from mypyc.common import BITMAP_BITS, ENV_ATTR_NAME, SELF_NAME, bitmap_name
+from mypyc.common import (
+    BITMAP_BITS,
+    ENV_ATTR_NAME,
+    GENERATOR_ATTRIBUTE_PREFIX,
+    SELF_NAME,
+    bitmap_name,
+)
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.ops import Call, GetAttr, SetAttr, Value
 from mypyc.ir.rtypes import RInstance, bitmap_rprimitive, object_rprimitive
@@ -60,7 +66,7 @@ class is generated, the function environment has not yet been
     return env_class
 
 
-def finalize_env_class(builder: IRBuilder) -> None:
+def finalize_env_class(builder: IRBuilder, prefix: str = "") -> None:
     """Generate, instantiate, and set up the environment of an environment class."""
     if not builder.fn_info.can_merge_generator_and_env_classes():
         instantiate_env_class(builder)
@@ -69,9 +75,9 @@ def finalize_env_class(builder: IRBuilder) -> None:
     # that were previously added to the environment with references to the function's
     # environment class.
     if builder.fn_info.is_nested:
-        add_args_to_env(builder, local=False, base=builder.fn_info.callable_class)
+        add_args_to_env(builder, local=False, base=builder.fn_info.callable_class, prefix=prefix)
     else:
-        add_args_to_env(builder, local=False, base=builder.fn_info)
+        add_args_to_env(builder, local=False, base=builder.fn_info, prefix=prefix)
 
 
 def instantiate_env_class(builder: IRBuilder) -> Value:
@@ -96,15 +102,15 @@ def instantiate_env_class(builder: IRBuilder) -> Value:
     return curr_env_reg
 
 
-def load_env_registers(builder: IRBuilder) -> None:
+def load_env_registers(builder: IRBuilder, prefix: str = "") -> None:
     """Load the registers for the current FuncItem being visited.
 
     Adds the arguments of the FuncItem to the environment. If the
     FuncItem is nested inside of another function, then this also
     loads all of the outer environments of the FuncItem into registers
     so that they can be used when accessing free variables.
     """
-    add_args_to_env(builder, local=True)
+    add_args_to_env(builder, local=True, prefix=prefix)
 
     fn_info = builder.fn_info
     fitem = fn_info.fitem
@@ -113,7 +119,7 @@ def load_env_registers(builder: IRBuilder) -> None:
         # If this is a FuncDef, then make sure to load the FuncDef into its own environment
         # class so that the function can be called recursively.
         if isinstance(fitem, FuncDef) and fn_info.add_nested_funcs_to_env:
-            setup_func_for_recursive_call(builder, fitem, fn_info.callable_class)
+            setup_func_for_recursive_call(builder, fitem, fn_info.callable_class, prefix=prefix)
 
 
 def load_outer_env(
@@ -134,8 +140,11 @@ def load_outer_env(
     assert isinstance(env.type, RInstance), f"{env} must be of type RInstance"
 
     for symbol, target in outer_env.items():
-        env.type.class_ir.attributes[symbol.name] = target.type
-        symbol_target = AssignmentTargetAttr(env, symbol.name)
+        attr_name = symbol.name
+        if isinstance(target, AssignmentTargetAttr):
+            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)
 
     return env
@@ -178,6 +187,7 @@ def add_args_to_env(
     local: bool = True,
     base: FuncInfo | ImplicitClass | None = None,
     reassign: bool = True,
+    prefix: str = "",
 ) -> None:
     fn_info = builder.fn_info
     args = fn_info.fitem.arguments
@@ -193,10 +203,12 @@ def add_args_to_env(
             if is_free_variable(builder, arg.variable) or fn_info.is_generator:
                 rtype = builder.type_to_rtype(arg.variable.type)
                 assert base is not None, "base cannot be None for adding nonlocal args"
-                builder.add_var_to_env_class(arg.variable, rtype, base, reassign=reassign)
+                builder.add_var_to_env_class(
+                    arg.variable, rtype, base, reassign=reassign, prefix=prefix
+                )
 
 
-def add_vars_to_env(builder: IRBuilder) -> None:
+def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None:
     """Add relevant local variables and nested functions to the environment class.
 
     Add all variables and functions that are declared/defined within current
@@ -216,7 +228,9 @@ def add_vars_to_env(builder: IRBuilder) -> None:
         for var in sorted(builder.free_variables[builder.fn_info.fitem], key=lambda x: x.name):
             if isinstance(var, Var):
                 rtype = builder.type_to_rtype(var.type)
-                builder.add_var_to_env_class(var, rtype, env_for_func, reassign=False)
+                builder.add_var_to_env_class(
+                    var, rtype, env_for_func, reassign=False, prefix=prefix
+                )
 
     if builder.fn_info.fitem in builder.encapsulating_funcs:
         for nested_fn in builder.encapsulating_funcs[builder.fn_info.fitem]:
@@ -226,12 +240,16 @@ def add_vars_to_env(builder: IRBuilder) -> None:
                 # the same name and signature across conditional blocks
                 # will generate different callable classes, so the callable
                 # class that gets instantiated must be generic.
+                if nested_fn.is_generator:
+                    prefix = GENERATOR_ATTRIBUTE_PREFIX
                 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
                 )
 
 
-def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: ImplicitClass) -> None:
+def setup_func_for_recursive_call(
+    builder: IRBuilder, fdef: FuncDef, base: ImplicitClass, prefix: str = ""
+) -> None:
     """Enable calling a nested function (with a callable class) recursively.
 
     Adds the instance of the callable class representing the given
@@ -241,7 +259,8 @@ def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: Impli
     """
     # First, set the attribute of the environment class so that GetAttr can be called on it.
     prev_env = builder.fn_infos[-2].env_class
-    prev_env.attributes[fdef.name] = builder.type_to_rtype(fdef.type)
+    attr_name = prefix + fdef.name
+    prev_env.attributes[attr_name] = builder.type_to_rtype(fdef.type)
 
     if isinstance(base, GeneratorClass):
         # If we are dealing with a generator class, then we need to first get the register
@@ -253,7 +272,7 @@ def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: Impli
 
     # Obtain the instance of the callable class representing the FuncDef, and add it to the
     # current environment.
-    val = builder.add(GetAttr(prev_env_reg, fdef.name, -1))
+    val = builder.add(GetAttr(prev_env_reg, attr_name, -1))
     target = builder.add_local_reg(fdef, object_rprimitive)
     builder.assign(target, val, -1)
 

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 (
@@ -68,14 +68,14 @@ def gen_generator_func(
 ) -> tuple[FuncIR, Value | None]:
     """Generate IR for generator function that returns generator object."""
     setup_generator_class(builder)
-    load_env_registers(builder)
+    load_env_registers(builder, prefix=GENERATOR_ATTRIBUTE_PREFIX)
     gen_arg_defaults(builder)
     if builder.fn_info.can_merge_generator_and_env_classes():
         gen = instantiate_generator_class(builder)
         builder.fn_info._curr_env_reg = gen
-        finalize_env_class(builder)
+        finalize_env_class(builder, prefix=GENERATOR_ATTRIBUTE_PREFIX)
     else:
-        finalize_env_class(builder)
+        finalize_env_class(builder, prefix=GENERATOR_ATTRIBUTE_PREFIX)
         gen = instantiate_generator_class(builder)
     builder.add(Return(gen))
 
@@ -104,11 +104,13 @@ class that implements the function (each function gets a separate class).
         and top_level
         and top_level.add_nested_funcs_to_env
     ):
-        setup_func_for_recursive_call(builder, fitem, builder.fn_info.generator_class)
+        setup_func_for_recursive_call(
+            builder, fitem, builder.fn_info.generator_class, prefix=GENERATOR_ATTRIBUTE_PREFIX
+        )
     create_switch_for_generator_class(builder)
     add_raise_exception_blocks_to_generator_class(builder, fitem.line)
 
-    add_vars_to_env(builder)
+    add_vars_to_env(builder, prefix=GENERATOR_ATTRIBUTE_PREFIX)
 
     builder.accept(fitem.body)
     builder.maybe_add_implicit_return()
@@ -429,7 +431,9 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None:
 
     # Add arguments from the original generator function to the
     # environment of the generator class.
-    add_args_to_env(builder, local=False, base=cls, reassign=False)
+    add_args_to_env(
+        builder, local=False, base=cls, reassign=False, prefix=GENERATOR_ATTRIBUTE_PREFIX
+    )
 
     # Set the next label register for the generator class.
     cls.next_label_reg = builder.read(cls.next_label_target, fitem.line)

mypyc/test-data/run-async.test

@@ -1208,3 +1208,29 @@ async def test_async_context_manager_exception_handling() -> None:
 
 [file asyncio/__init__.pyi]
 async def sleep(t: float) -> None: ...
+
+[case testCallableArgWithSameNameAsHelperMethod]
+import asyncio
+from typing import Awaitable, Callable
+
+
+MyCallable = Callable[[int, int], Awaitable[int]]
+
+async def add(a: int, b: int) -> int:
+    return a + b
+
+async def await_send(send: MyCallable) -> int:
+    return await send(1, 2)
+
+async def await_throw(throw: MyCallable) -> int:
+    return await throw(3, 4)
+
+async def tests() -> None:
+    assert await await_send(add) == 3
+    assert await await_throw(add) == 7
+
+def test_callable_arg_same_name_as_helper() -> None:
+    asyncio.run(tests())
+
+[file asyncio/__init__.pyi]
+def run(x: object) -> object: ...

mypyc/test-data/run-generators.test

@@ -272,8 +272,17 @@ def call_nested_decorated(x: int) -> list[int]:
         a.append(x)
     return a
 
+def call_nested_recursive(x: int) -> Iterator:
+    def recursive(x: int) -> Iterator:
+        if x > 0:
+            yield from recursive(x - 1)
+        yield x
+
+    yield from recursive(x)
+
 def test_call_nested_generator_in_function() -> None:
     assert call_nested_decorated(5) == [5, 15]
+    assert list(call_nested_recursive(5)) == [0, 1, 2, 3, 4, 5]
 
 [case testYieldThrow]
 from typing import Generator, Iterable, Any, Union
@@ -871,3 +880,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())