[mypyc] Fix undefined attribute in nested coroutines

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: ...