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

[case testAsyncDefPass]

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

[case testAsyncDefReturn]

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

[case testAsyncDefMissingReturn]
# flags: --warn-no-return
async def f() -> int:
    make_this_not_trivial = 1
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-full.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-full.pyi]
[out]
main:4: error: Return value expected

[case testAwaitCoroutine]

async def f() -> int:
    x = await f()
    reveal_type(x)  # E: Revealed type is 'builtins.int*'
    return x
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-full.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-full.pyi]
[out]
main:6: error: 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-full.pyi]
[out]
main:6: error: 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-full.pyi]
[out]
main:5: error: Argument 1 to "f" has incompatible type "T"; expected "int"
main:6: error: 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-full.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-full.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-full.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-full.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-full.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)  # E: Revealed type is 'builtins.int*'
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-full.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-full.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)  # E: Revealed type is 'Union[builtins.int, builtins.str]'
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-full.pyi]

[case testAsyncForComprehension]
# flags: --fast-parser --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)  # E: Revealed type is 'builtins.list[builtins.int*]'
    lst2 = [i async for i in asyncify(obj) for j in obj]
    reveal_type(lst2)  # E: Revealed type is 'builtins.list[builtins.int*]'

async def setcomp(obj: Iterable[int]):
    lst = {i async for i in asyncify(obj)}
    reveal_type(lst)  # E: 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)  # E: 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)  # E: Revealed type is 'typing.AsyncGenerator[builtins.int*, None]'
    lst2 = (i async for i in asyncify(obj) for i in obj)
    reveal_type(lst2)  # E: Revealed type is 'typing.AsyncGenerator[builtins.int*, None]'

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

[case testAsyncForComprehensionErrors]
# flags: --fast-parser --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-full.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)  # E: Revealed type is 'builtins.int*'
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-full.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-full.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-full.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-full.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-full.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-full.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-full.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-full.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)  # E: Revealed type is 'builtins.int'
    async with C() as y:
        reveal_type(y)  # E: Revealed type is 'builtins.int'
    async for z in I():
        reveal_type(z)  # E: Revealed type is 'builtins.int'
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-full.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)  # E: 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-full.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)  # E: 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)  # E: Revealed type is 'def () -> typing.AsyncGenerator[builtins.int, None]'
reveal_type(g())  # E: Revealed type is 'typing.AsyncGenerator[builtins.int, None]'

async def h() -> None:
    async for item in g():
        reveal_type(item)  # E: 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-full.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)  # E: Revealed type is 'builtins.int*'

[builtins fixtures/dict.pyi]
[typing fixtures/typing-full.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__())  # E: Revealed type is 'typing.AsyncGenerator[builtins.int*, None]'

    reveal_type(await gen.__anext__())  # E: Revealed type is 'builtins.int*'

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

[case testAsyncGeneratorAsend]
# flags: --fast-parser --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)  # E: 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'))  # E: Revealed type is 'builtins.int*'

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

[case testAsyncGeneratorAthrow]
# flags: --fast-parser --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)  # E: Revealed type is 'builtins.str*'
    reveal_type(await g.athrow(BaseException))  # E: Revealed type is 'builtins.str*'

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

[case testAsyncGeneratorNoSyncIteration]
# flags: --fast-parser --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-full.pyi]

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

[case testAsyncGeneratorNoYieldFrom]
# flags: --fast-parser --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-full.pyi]

[case testAsyncGeneratorNoReturnWithValue]
# flags: --fast-parser --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-full.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-full.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-full.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 type annotation
    yield 0
[builtins fixtures/async_await.pyi]
[typing fixtures/typing-full.pyi]
[out]
