Skip to content

Commit

Permalink
Fix nested async functions when using TypeVar value restriction (#14705)
Browse files Browse the repository at this point in the history
Propagate additional function flags in TransformVisitor.

Work on #14706.
  • Loading branch information
JukkaL authored Apr 24, 2023
1 parent 186432d commit bdac4bc
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 0 deletions.
3 changes: 3 additions & 0 deletions mypy/treetransform.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
42 changes: 42 additions & 0 deletions test-data/unit/check-async-await.test
Original file line number Diff line number Diff line change
Expand Up @@ -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]

0 comments on commit bdac4bc

Please sign in to comment.