Fix nested async functions when using TypeVar value restriction (#14705)

Propagate additional function flags in TransformVisitor.

Work on #14706.

Author: JukkaL

mypy/treetransform.py

@@ -233,6 +233,9 @@ def copy_function_attributes(self, new: FuncItem, original: FuncItem) -> None:
         new.max_pos = original.max_pos
         new.is_overload = original.is_overload
         new.is_generator = original.is_generator
+        new.is_coroutine = original.is_coroutine
+        new.is_async_generator = original.is_async_generator
+        new.is_awaitable_coroutine = original.is_awaitable_coroutine
         new.line = original.line
 
     def visit_overloaded_func_def(self, node: OverloadedFuncDef) -> OverloadedFuncDef:

test-data/unit/check-async-await.test

@@ -959,3 +959,45 @@ async def good() -> None:
     y = [await foo(x) for x in [1, 2, 3]]  # OK
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]
+
+[case testNestedAsyncFunctionAndTypeVarAvalues]
+from typing import TypeVar
+
+T = TypeVar('T', int, str)
+
+def f(x: T) -> None:
+    async def g() -> T:
+        return x
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]
+
+[case testNestedAsyncGeneratorAndTypeVarAvalues]
+from typing import AsyncGenerator, TypeVar
+
+T = TypeVar('T', int, str)
+
+def f(x: T) -> None:
+    async def g() -> AsyncGenerator[T, None]:
+        yield x
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]
+
+[case testNestedDecoratedCoroutineAndTypeVarValues]
+from typing import Generator, TypeVar
+from types import coroutine
+
+T = TypeVar('T', int, str)
+
+def f(x: T) -> None:
+    @coroutine
+    def inner() -> Generator[T, None, None]:
+        yield x
+    reveal_type(inner)  # N: Revealed type is "def () -> typing.AwaitableGenerator[builtins.int, None, None, typing.Generator[builtins.int, None, None]]" \
+        # N: Revealed type is "def () -> typing.AwaitableGenerator[builtins.str, None, None, typing.Generator[builtins.str, None, None]]"
+
+@coroutine
+def coro() -> Generator[int, None, None]:
+    yield 1
+reveal_type(coro)  # N: Revealed type is "def () -> typing.AwaitableGenerator[builtins.int, None, None, typing.Generator[builtins.int, None, None]]"
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]