check-async-await.test

-- Tests for async def and await (PEP 492)
-- ---------------------------------------

[case testAsyncDefPass]

async def f() -> int:
    pass
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncDefReturn]

async def f() -> int:
    return 0
reveal_type(f())  # N: Revealed type is "typing.Coroutine[Any, Any, builtins.int]"
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncDefMissingReturn]
# flags: --warn-no-return
async def f() -> int:
    make_this_not_trivial = 1
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]
main:2: error: Missing return statement

[case testAsyncDefReturnWithoutValue]

async def f() -> int:
    make_this_not_trivial = 1
    return
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]
main:4: error: Return value expected

[case testAwaitCoroutine]

async def f() -> int:
    x = await f()
    reveal_type(x)  # N: Revealed type is "builtins.int*"
    return x
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]

[case testAwaitDefaultContext]

from typing import TypeVar
T = TypeVar('T')
async def f(x: T) -> T:
    y = await f(x)
    reveal_type(y)
    return y
[typing fixtures/typing-async.pyi]
[out]
main:6: note: Revealed type is "T`-1"

[case testAwaitAnyContext]

from typing import Any, TypeVar
T = TypeVar('T')
async def f(x: T) -> T:
    y = await f(x)  # type: Any
    reveal_type(y)
    return y
[typing fixtures/typing-async.pyi]
[out]
main:6: note: Revealed type is "Any"

[case testAwaitExplicitContext]

from typing import TypeVar
T = TypeVar('T')
async def f(x: T) -> T:
    y = await f(x)  # type: int
    reveal_type(y)
    return x
[typing fixtures/typing-async.pyi]
[out]
main:5: error: Argument 1 to "f" has incompatible type "T"; expected "int"
main:6: note: Revealed type is "builtins.int"

[case testAwaitGeneratorError]

from typing import Any, Generator
def g() -> Generator[int, None, str]:
    yield 0
    return ''
async def f() -> int:
    x = await g()
    return x
[typing fixtures/typing-async.pyi]
[out]
main:7: error: Incompatible types in "await" (actual type "Generator[int, None, str]", expected type "Awaitable[Any]")

[case testAwaitIteratorError]

from typing import Any, Iterator
def g() -> Iterator[Any]:
    yield
async def f() -> int:
    x = await g()
    return x
[typing fixtures/typing-async.pyi]
[out]
main:6: error: Incompatible types in "await" (actual type "Iterator[Any]", expected type "Awaitable[Any]")

[case testAwaitArgumentError]

def g() -> int:
    return 0
async def f() -> int:
    x = await g()
    return x
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]
main:5: error: Incompatible types in "await" (actual type "int", expected type "Awaitable[Any]")

[case testAwaitResultError]

async def g() -> int:
    return 0
async def f() -> str:
    x = await g()  # type: str
    return x
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]
main:5: error: Incompatible types in assignment (expression has type "int", variable has type "str")

[case testAwaitReturnError]

async def g() -> int:
    return 0
async def f() -> str:
    x = await g()
    return x
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]
main:6: error: Incompatible return value type (got "int", expected "str")

[case testAsyncFor]

from typing import AsyncIterator
class C(AsyncIterator[int]):
    async def __anext__(self) -> int: return 0
async def f() -> None:
    async for x in C():
        reveal_type(x)  # N: Revealed type is "builtins.int*"
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncForError]

from typing import AsyncIterator
async def f() -> None:
    async for x in [1]:
        pass
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]
main:4: error: "List[int]" has no attribute "__aiter__" (not async iterable)

[case testAsyncForTypeComments]

from typing import AsyncIterator, Union
class C(AsyncIterator[int]):
    async def __anext__(self) -> int: return 0
async def f() -> None:
    async for x in C():  # type: str  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
        pass

    async for y in C():  # type: int
        pass

    async for z in C():  # type: Union[int, str]
        reveal_type(z)  # N: Revealed type is "Union[builtins.int, builtins.str]"
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncForComprehension]
# flags: --python-version 3.6
from typing import Generic, Iterable, TypeVar, AsyncIterator, Tuple

T = TypeVar('T')

class asyncify(Generic[T], AsyncIterator[T]):
    def __init__(self, iterable: Iterable[T]) -> None:
        self.iterable = iter(iterable)
    def __aiter__(self) -> AsyncIterator[T]:
        return self
    async def __anext__(self) -> T:
        try:
            return next(self.iterable)
        except StopIteration:
            raise StopAsyncIteration

async def listcomp(obj: Iterable[int]):
    lst = [i async for i in asyncify(obj)]
    reveal_type(lst)  # N: Revealed type is "builtins.list[builtins.int*]"
    lst2 = [i async for i in asyncify(obj) for j in obj]
    reveal_type(lst2)  # N: Revealed type is "builtins.list[builtins.int*]"

async def setcomp(obj: Iterable[int]):
    lst = {i async for i in asyncify(obj)}
    reveal_type(lst)  # N: Revealed type is "builtins.set[builtins.int*]"

async def dictcomp(obj: Iterable[Tuple[int, str]]):
    lst = {i: j async for i, j in asyncify(obj)}
    reveal_type(lst)  # N: Revealed type is "builtins.dict[builtins.int*, builtins.str*]"

async def generatorexp(obj: Iterable[int]):
    lst = (i async for i in asyncify(obj))
    reveal_type(lst)  # N: Revealed type is "typing.AsyncGenerator[builtins.int*, None]"
    lst2 = (i async for i in asyncify(obj) for i in obj)
    reveal_type(lst2)  # N: Revealed type is "typing.AsyncGenerator[builtins.int*, None]"

[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncForComprehensionErrors]
# flags: --python-version 3.6
from typing import Generic, Iterable, TypeVar, AsyncIterator, Tuple

T = TypeVar('T')

class asyncify(Generic[T], AsyncIterator[T]):
    def __init__(self, iterable: Iterable[T]) -> None:
        self.iterable = iter(iterable)
    def __aiter__(self) -> AsyncIterator[T]:
        return self
    async def __anext__(self) -> T:
        try:
            return next(self.iterable)
        except StopIteration:
            raise StopAsyncIteration

async def wrong_iterable(obj: Iterable[int]):
    [i async for i in obj]
    [i for i in asyncify(obj)]
    {i: i async for i in obj}
    {i: i for i in asyncify(obj)}

[out]
main:18: error: "Iterable[int]" has no attribute "__aiter__" (not async iterable)
main:19: error: "asyncify[int]" has no attribute "__iter__"; maybe "__aiter__"? (not iterable)
main:20: error: "Iterable[int]" has no attribute "__aiter__" (not async iterable)
main:21: error: "asyncify[int]" has no attribute "__iter__"; maybe "__aiter__"? (not iterable)
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncWith]

class C:
    async def __aenter__(self) -> int: pass
    async def __aexit__(self, x, y, z) -> None: pass
async def f() -> None:
    async with C() as x:
        reveal_type(x)  # N: Revealed type is "builtins.int*"
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]


[case testAsyncWithError]

class C:
    def __enter__(self) -> int: pass
    def __exit__(self, x, y, z) -> None: pass
async def f() -> None:
    async with C() as x:
        pass
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]
main:6: error: "C" has no attribute "__aenter__"; maybe "__enter__"?
main:6: error: "C" has no attribute "__aexit__"; maybe "__exit__"?

[case testAsyncWithErrorBadAenter]

class C:
    def __aenter__(self) -> int: pass
    async def __aexit__(self, x, y, z) -> None: pass
async def f() -> None:
    async with C() as x:  # E: Incompatible types in "async with" for "__aenter__" (actual type "int", expected type "Awaitable[Any]")
        pass
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncWithErrorBadAenter2]

class C:
    def __aenter__(self) -> None: pass
    async def __aexit__(self, x, y, z) -> None: pass
async def f() -> None:
    async with C() as x:  # E: "None" has no attribute "__await__"
        pass
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncWithErrorBadAexit]

class C:
    async def __aenter__(self) -> int: pass
    def __aexit__(self, x, y, z) -> int: pass
async def f() -> None:
    async with C() as x: # E: Incompatible types in "async with" for "__aexit__" (actual type "int", expected type "Awaitable[Any]")
        pass
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncWithErrorBadAexit2]

class C:
    async def __aenter__(self) -> int: pass
    def __aexit__(self, x, y, z) -> None: pass
async def f() -> None:
    async with C() as x: # E: "None" has no attribute "__await__"
        pass
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncWithTypeComments]

class C:
    async def __aenter__(self) -> int: pass
    async def __aexit__(self, x, y, z) -> None: pass
async def f() -> None:
    async with C() as x:  # type: int
        pass

    async with C() as y, C() as z:  # type: str, int  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
        pass

    async with C() as a:  # type: int, int  # E: Syntax error in type annotation # N: Suggestion: Use Tuple[T1, ..., Tn] instead of (T1, ..., Tn)
        pass
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testNoYieldInAsyncDef]
# flags: --python-version 3.5

async def f():
    yield None  # E: "yield" in async function
async def g():
    yield  # E: "yield" in async function
async def h():
    x = yield  # E: "yield" in async function
[builtins fixtures/async_await.pyi]

[case testNoYieldFromInAsyncDef]

async def f():
    yield from []
async def g():
    x = yield from []
[builtins fixtures/async_await.pyi]
[out]
main:3: error: "yield from" in async function
main:5: error: "yield from" in async function

[case testNoAsyncDefInPY2_python2]

async def f():  # E: invalid syntax
    pass

[case testYieldFromNoAwaitable]

from typing import Any, Generator
async def f() -> str:
    return ''
def g() -> Generator[Any, None, str]:
    x = yield from f()
    return x
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]
main:6: error: "yield from" can't be applied to "Coroutine[Any, Any, str]"

[case testAwaitableSubclass]

from typing import Any, AsyncIterator, Awaitable, Generator
class A(Awaitable[int]):
    def __await__(self) -> Generator[Any, None, int]:
        yield
        return 0
class C:
    def __aenter__(self) -> A:
        return A()
    def __aexit__(self, *a) -> A:
        return A()
class I(AsyncIterator[int]):
    def __aiter__(self) -> 'I':
        return self
    def __anext__(self) -> A:
        return A()
async def main() -> None:
    x = await A()
    reveal_type(x)  # N: Revealed type is "builtins.int"
    async with C() as y:
        reveal_type(y)  # N: Revealed type is "builtins.int"
    async for z in I():
        reveal_type(z)  # N: Revealed type is "builtins.int"
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testYieldTypeCheckInDecoratedCoroutine]

from typing import Generator
from types import coroutine
@coroutine
def f() -> Generator[int, str, int]:
    x = yield 0
    x = yield ''  # E: Incompatible types in "yield" (actual type "str", expected type "int")
    reveal_type(x)  # N: Revealed type is "builtins.str"
    if x:
        return 0
    else:
        return ''  # E: Incompatible return value type (got "str", expected "int")
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

-- Async generators (PEP 525), some test cases adapted from the PEP text
-- ---------------------------------------------------------------------

[case testAsyncGenerator]
# flags: --python-version 3.6
from typing import AsyncGenerator, Generator

async def f() -> int:
    return 42

async def g() -> AsyncGenerator[int, None]:
    value = await f()
    reveal_type(value)  # N: Revealed type is "builtins.int*"
    yield value

    yield 'not an int'  # E: Incompatible types in "yield" (actual type "str", expected type "int")
    # return without a value is fine
    return
reveal_type(g)  # N: Revealed type is "def () -> typing.AsyncGenerator[builtins.int, None]"
reveal_type(g())  # N: Revealed type is "typing.AsyncGenerator[builtins.int, None]"

async def h() -> None:
    async for item in g():
        reveal_type(item)  # N: Revealed type is "builtins.int*"

async def wrong_return() -> Generator[int, None, None]:  # E: The return type of an async generator function should be "AsyncGenerator" or one of its supertypes
    yield 3

[builtins fixtures/dict.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncGeneratorReturnIterator]
# flags: --python-version 3.6
from typing import AsyncIterator

async def gen() -> AsyncIterator[int]:
    yield 3

    yield 'not an int'  # E: Incompatible types in "yield" (actual type "str", expected type "int")

async def use_gen() -> None:
    async for item in gen():
        reveal_type(item)  # N: Revealed type is "builtins.int*"

[builtins fixtures/dict.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncGeneratorManualIter]
# flags: --python-version 3.6
from typing import AsyncGenerator

async def genfunc() -> AsyncGenerator[int, None]:
    yield 1
    yield 2

async def user() -> None:
    gen = genfunc()

    reveal_type(gen.__aiter__())  # N: Revealed type is "typing.AsyncGenerator[builtins.int*, None]"

    reveal_type(await gen.__anext__())  # N: Revealed type is "builtins.int*"

[builtins fixtures/dict.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncGeneratorAsend]
# flags: --python-version 3.6
from typing import AsyncGenerator

async def f() -> None:
    pass

async def gen() -> AsyncGenerator[int, str]:
    await f()
    v = yield 42
    reveal_type(v)  # N: Revealed type is "builtins.str"
    await f()

async def h() -> None:
    g = gen()
    await g.asend(())  # E: Argument 1 to "asend" of "AsyncGenerator" has incompatible type "Tuple[]"; expected "str"
    reveal_type(await g.asend('hello'))  # N: Revealed type is "builtins.int*"

[builtins fixtures/dict.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncGeneratorAthrow]
# flags: --python-version 3.6
from typing import AsyncGenerator

async def gen() -> AsyncGenerator[str, int]:
    try:
        yield 'hello'
    except BaseException:
        yield 'world'

async def h() -> None:
    g = gen()
    v = await g.asend(1)
    reveal_type(v)  # N: Revealed type is "builtins.str*"
    reveal_type(await g.athrow(BaseException))  # N: Revealed type is "builtins.str*"

[builtins fixtures/dict.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncGeneratorNoSyncIteration]
# flags: --python-version 3.6
from typing import AsyncGenerator

async def gen() -> AsyncGenerator[int, None]:
    for i in [1, 2, 3]:
        yield i

def h() -> None:
    for i in gen():
        pass

[builtins fixtures/dict.pyi]
[typing fixtures/typing-async.pyi]

[out]
main:9: error: "AsyncGenerator[int, None]" has no attribute "__iter__"; maybe "__aiter__"? (not iterable)

[case testAsyncGeneratorNoYieldFrom]
# flags: --python-version 3.6
from typing import AsyncGenerator

async def f() -> AsyncGenerator[int, None]:
    pass

async def gen() -> AsyncGenerator[int, None]:
    yield from f()  # E: "yield from" in async function

[builtins fixtures/dict.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncGeneratorNoReturnWithValue]
# flags: --python-version 3.6
from typing import AsyncGenerator

async def return_int() -> AsyncGenerator[int, None]:
    yield 1
    return 42  # E: "return" with value in async generator is not allowed

async def return_none() -> AsyncGenerator[int, None]:
    yield 1
    return None  # E: "return" with value in async generator is not allowed

def f() -> None:
    return

async def return_f() -> AsyncGenerator[int, None]:
    yield 1
    return f()  # E: "return" with value in async generator is not allowed

[builtins fixtures/dict.pyi]
[typing fixtures/typing-async.pyi]

-- The full matrix of coroutine compatibility
-- ------------------------------------------

[case testFullCoroutineMatrix]

from typing import Any, AsyncIterator, Awaitable, Generator, Iterator
from types import coroutine

# The various things you might try to use in `await` or `yield from`.

def plain_generator() -> Generator[str, None, int]:
    yield 'a'
    return 1

async def plain_coroutine() -> int:
    return 1

@coroutine
def decorated_generator() -> Generator[str, None, int]:
    yield 'a'
    return 1

@coroutine
async def decorated_coroutine() -> int:
    return 1

class It(Iterator[str]):
    def __iter__(self) -> 'It':
        return self
    def __next__(self) -> str:
        return 'a'

def other_iterator() -> It:
    return It()

class Aw(Awaitable[int]):
    def __await__(self) -> Generator[str, Any, int]:
        yield 'a'
        return 1

def other_coroutine() -> Aw:
    return Aw()

# The various contexts in which `await` or `yield from` might occur.

def plain_host_generator() -> Generator[str, None, None]:
    yield 'a'
    x = 0
    x = yield from plain_generator()
    x = yield from plain_coroutine()  # E: "yield from" can't be applied to "Coroutine[Any, Any, int]"
    x = yield from decorated_generator()
    x = yield from decorated_coroutine()  # E: "yield from" can't be applied to "AwaitableGenerator[Any, Any, int, Coroutine[Any, Any, int]]"
    x = yield from other_iterator()
    x = yield from other_coroutine()  # E: "yield from" can't be applied to "Aw"

async def plain_host_coroutine() -> None:
    x = 0
    x = await plain_generator()  # E: Incompatible types in "await" (actual type "Generator[str, None, int]", expected type "Awaitable[Any]")
    x = await plain_coroutine()
    x = await decorated_generator()
    x = await decorated_coroutine()
    x = await other_iterator()  # E: Incompatible types in "await" (actual type "It", expected type "Awaitable[Any]")
    x = await other_coroutine()

@coroutine
def decorated_host_generator() -> Generator[str, None, None]:
    yield 'a'
    x = 0
    x = yield from plain_generator()
    x = yield from plain_coroutine()
    x = yield from decorated_generator()
    x = yield from decorated_coroutine()
    x = yield from other_iterator()
    x = yield from other_coroutine()  # E: "yield from" can't be applied to "Aw"

@coroutine
async def decorated_host_coroutine() -> None:
    x = 0
    x = await plain_generator()  # E: Incompatible types in "await" (actual type "Generator[str, None, int]", expected type "Awaitable[Any]")
    x = await plain_coroutine()
    x = await decorated_generator()
    x = await decorated_coroutine()
    x = await other_iterator()  # E: Incompatible types in "await" (actual type "It", expected type "Awaitable[Any]")
    x = await other_coroutine()

[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]

[case testAsyncGenDisallowUntyped]
# flags: --disallow-untyped-defs
# These should not crash
from typing import AsyncGenerator, Any

async def f() -> AsyncGenerator[int, None]:
    yield 0

async def g() -> AsyncGenerator[Any, None]:
    yield 0
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]

[case testAsyncGenDisallowUntypedTriggers]
# flags: --disallow-untyped-defs
from typing import AsyncGenerator, Any

async def f() -> AsyncGenerator[Any, Any]:
    yield None

async def h() -> Any:
    yield 0

async def g():  # E: Function is missing a return type annotation
    yield 0
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]
[out]

[case testAsyncOverloadedFunction]
from typing import overload

@overload
async def f(x: int) -> int: ...
@overload
async def f(x: str) -> str: ...
async def f(x):
    pass

reveal_type(f) # N: Revealed type is "Overload(def (x: builtins.int) -> typing.Coroutine[Any, Any, builtins.int], def (x: builtins.str) -> typing.Coroutine[Any, Any, builtins.str])"
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncForwardRefInBody]

async def f() -> None:
    forwardref: C
    class C: pass

def dec(x): pass

@dec
async def g() -> None:
    forwardref: C
    class C: pass

reveal_type(f) # N: Revealed type is "def () -> typing.Coroutine[Any, Any, None]"
reveal_type(g) # N: Revealed type is "Any"
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncDeferredAnalysis]

def t() -> None:
    async def f() -> int:
        return 1

    def g() -> int:
        return next(iter(x))

    x = [1]

[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAsyncWithInGenericClass]
from typing import Generic, AsyncContextManager, TypeVar

T = TypeVar('T', str, int)

class Foo(Generic[T]):
    async def foo(self, manager: AsyncContextManager):
        async with manager:
            pass

[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]

[case testAwaitOverloadSpecialCase]
from typing import Any, Awaitable, Iterable, overload, Tuple, List, TypeVar, Generic

T = TypeVar("T")
FT = TypeVar("FT", bound='Future[Any]')

class Future(Awaitable[T], Iterable[T]):
    pass

class Task(Future[T]):
    pass

@overload
def wait(fs: Iterable[FT]) -> Future[Tuple[List[FT], List[FT]]]: ...  \
    # E: Overloaded function signatures 1 and 2 overlap with incompatible return types
@overload
def wait(fs: Iterable[Awaitable[T]]) -> Future[Tuple[List[Task[T]], List[Task[T]]]]: ...
def wait(fs: Any) -> Any:
    pass

async def imprecise1(futures: Iterable[Task[Any]]) -> None:
    done: Any
    pending: Any
    done, pending = await wait(futures)
    reveal_type(done)  # N: Revealed type is "Any"

async def imprecise2(futures: Iterable[Awaitable[Any]]) -> None:
    done, pending = await wait(futures)
    reveal_type(done)  # N: Revealed type is "builtins.list[__main__.Task[Any]]"

async def precise1(futures: Iterable[Future[int]]) -> None:
    done, pending = await wait(futures)
    reveal_type(done)  # N: Revealed type is "builtins.list[__main__.Future[builtins.int]]"

async def precise2(futures: Iterable[Awaitable[int]]) -> None:
    done, pending = await wait(futures)
    reveal_type(done)  # N: Revealed type is "builtins.list[__main__.Task[builtins.int]]"

[builtins fixtures/async_await.pyi]
[typing fixtures/typing-async.pyi]