diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py index 3f5a05386e7a7..164765dd8db42 100644 --- a/mypyc/irbuild/builder.py +++ b/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, diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py index 2a9cc50d9ccf1..a693405178631 100644 --- a/mypyc/irbuild/env_class.py +++ b/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 diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index f0320e60ee1a4..3b1bd9def18cf 100644 --- a/mypyc/test-data/run-async.test +++ b/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: ...