[mypyc] Fix undefined attribute in nested coroutines (#20654)

When creating a callable class for a generator, mypyc creates an
attribute in its env class that points at the compiled function so it
can be called recursively
https://github.com/python/mypy/blob/544b97ec296e16e4c96a42a1b81d5dcd841cf232/mypyc/irbuild/env_class.py#L263

The attribute is prefixed by `__mypyc_generator_attribute__` in case the
function name clashes with an attribute generated by mypyc internally.
However, in other places the attribute is currently not prefixed for
async functions which leads to mypyc generating both the prefixed and
unprefixed versions, eg. for a function `wrapper_async` the env class
will have both `_wrapper_async` and
`___mypyc_generator_attribute__wrapper_async` attributes with the second
one potentially undefined. This ends up in a runtime crash when another
nested function is called from a nested coroutine.

The issue is fixed by prefixing the attribute for coroutines in addition
to generators so that only the prefixed attribute is in the env class.

Author: p-sawicki

mypyc/irbuild/builder.py

@@ -657,7 +657,7 @@ def get_assignment_target(
                     # refers to the newly defined variable in that environment class. Add the
                     # target to the table containing class environment variables, as well as the
                     # current environment.
-                    if self.fn_info.is_generator:
+                    if self.fn_info.is_generator or self.fn_info.is_coroutine:
                         return self.add_var_to_env_class(
                             symbol,
                             reg_type,

mypyc/irbuild/env_class.py

@@ -200,7 +200,11 @@ def add_args_to_env(
             builder.add_local_reg(Var(bitmap_name(i)), bitmap_rprimitive, is_arg=True)
     else:
         for arg in args:
-            if is_free_variable(builder, arg.variable) or fn_info.is_generator:
+            if (
+                is_free_variable(builder, arg.variable)
+                or fn_info.is_generator
+                or fn_info.is_coroutine
+            ):
                 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(
@@ -240,7 +244,7 @@ def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> 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:
+                if nested_fn.is_generator or nested_fn.is_coroutine:
                     prefix = GENERATOR_ATTRIBUTE_PREFIX
                 builder.add_var_to_env_class(
                     nested_fn, object_rprimitive, env_for_func, reassign=False, prefix=prefix

mypyc/test-data/run-async.test

@@ -1799,3 +1799,49 @@ for i in range(10):
 from typing import Any, Generator
 
 def run(x: object) -> object: ...
+
+[case testNestedCoroutineCallsAnotherNestedFunction]
+import asyncio
+import functools
+import inspect
+from typing import Any, Callable, TypeVar, cast
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def mult(x: int) -> Callable[[F], F]:
+    def decorate(fn: F) -> F:
+        def get_multiplier() -> int:
+            return x
+
+        if inspect.iscoroutinefunction(fn):
+            @functools.wraps(fn)
+            async def wrapper_async(*args, **kwargs) -> Any:
+                return get_multiplier() * await fn(*args, **kwargs)
+            wrapper = wrapper_async
+        else:
+            @functools.wraps(fn)
+            def wrapper_non_async(*args, **kwargs) -> Any:
+                return get_multiplier() * fn(*args, **kwargs)
+            wrapper = wrapper_non_async
+
+        return cast(F, wrapper)
+
+    return decorate
+
+@mult(3)
+def identity(x: int):
+    return x
+
+@mult(5)
+async def async_identity(x: int):
+    return x
+
+def test_nested_coroutine_calls_another_nested_function():
+    assert identity(1) == 3
+    assert asyncio.run(async_identity(2)) == 10
+
+[file asyncio/__init__.pyi]
+from typing import Any, Generator
+
+def run(x: object) -> object: ...