Design for using async+await with asyncio

[gvanrossum]

I'm trying something new here, pouring my thoughts into the issue tracker to get some clarity.

I hope that we'll soon have syntactic support for PEP 492, in particular async def and await (#1808). (Yes, it also supports async for and async with, but these aren't very interesting so I won't mention them.)

I've also made some improvements to the asyncio stubs (python/typeshed#373, merged already) and created a fully annotated example, crawl.py (#1878). But these do not use async def or await; they use generators annotated with @asyncio.coroutine and yield from to wait for coroutines and Futures.

The next step is supporting asyncio code that uses async def and await, and here I am a little stuck. On the one hand, PEP 492 draws a clear distinction between generators and native coroutines (the latter being the ones defined using async def). You can't use a generator with await, and you can't use a native coroutine with yield from. On the other hand, PEP 492 also allows an escape clause: by using the decorator @types.coroutine you can mark a generator as a coroutine, and then await accepts it. The @asyncio.coroutine decorator calls this decorator, in addition to doing other asyncio-specific things (some of which are only apparent in debug mode).

I would like mypy to model all this as well as possible, so that mypy will give a useful error message when you try to use a generator with await or a coroutine with yield from. But it should also understand that a generator decorated with @types.coroutine or @asyncio.coroutine is acceptable for await as well as for yield from (in the latter case it would be even better if it could also keep track of whether the generator containing the yield from is decorated with @asyncio.coroutine).

The final piece of the puzzle (or rather, contribution to the problem) is that it's not easy to specify the type for a decorator that modifies the signature of the function it wraps. What I wish I could do is to declare that the @asyncio.coroutine decorator takes a function with arguments AA and return type Generator[Any, Any, R] and returns a function with arguments AA and return type Awaitable[R] (which is what await requires -- the return value of await is then of type R). And I'd like the decorator to be overloaded so that it also handles a few other cases.

But neither PEP 484 nor mypy supports this "signature algebra", so the best I can think of is to special-case the two decorators (@types.coroutine and @asyncio.coroutine) in mypy, in a similar fashion as how it deals with some other built-in decorators (@property, @classmethod and probably a few more).

[To be continued]

[gvanrossum]

Focusing on asyncio, there are now two equivalent ways of waiting for things:

• yield from, inside a generator (def).
• await, inside a native coroutine (async def).
• loop.run_until_complete(), assuming the event loop isn't already running.

In the first two cases the function should also be decorated with @asyncio.coroutine.

There are also three kinds of things that are worth waiting:

• asyncio.Future instances.
• coroutine objects (returned by functions defined with async def).
• generator objects (returned by generator functions).

In the last two cases the function should also be decorated with @asyncio.coroutine.

Thinking aloud, I'm going to pretend that await and yield from are functions with a signature. We can think of their current definitions like this:

from typing import Any, Awaitable, Generator
T = TypeVar('T')
def await_(arg: Awaitable[T]) -> T: ...
def yield_from_(arg: Generator[Any, None, T]) -> T: ...

To this list we can add the current definition of run_until_complete():

class AbstractEventLoop:
    def run_until_complete(self,
            future: Union[Future[T], Awaitable[T], Generator[Any, Any, T]]) -> T: ...

Next we look at the definitions. When we define a generator function, the return type must explicitly use the Generator type, like this:

def count(n: int) -> Generator[int, None, str]:
    for i in range(n):
        yield i
    return 'ok'

Note that the string 'ok' is returned via the StopIteration exception, and can be extracted by catching that exception and using its value attribute. Or, more idiomatically, you can use yield from inside another generator, and the "return value" of yield from will correspond to the return value of the generator, while the values yielded by the inner generator will be yielded by the outer generator:

def outer() -> Generator[int, None, None]:
    x = yield from count(5)
    print(x)

(This is just a summary of the behavior defined by PEP 380, which introduced yield from as well as the meaning of returning a values from a generator. Note that the first type parameter to Generator corresponds to the type of the values yielded using a yield statement/expression, and the third corresponds to the values returned using return. The middle parameter corresponds to the type of value sent into the generator using .send(value) -- this is the type returned by yield.)

For simple generators we can also use Iterator (or even Iterable):

from typing import Iterator
def count(n: int) -> Iterator[int]:
    for i in range(n):
        yield i

Now let's use an example with yield from:

def count_repeatedly(n: int, repeats: int) -> Iterator[int]:
    for i in range(repeats):
        yield from count(n)

On the other hand, PEP 492's native coroutines use a more concise notation:

async def my_coro(a: int) -> int:
    await another_coro()
    return 42

This is the moral equivalent of the following more cumbersome version using generators:

def my_coro_as_gen(a: int) -> Generator[Any, None, int]:
    yield from another_coro_as_gen()
    return 42

Now consider the types revealed by mypy:

reveal_type(my_coro(0))
reveal_type(my_coro_as_gen(0))

This will reveal the types to be as follows:

• Awaitable[int]
• Generator[Any, None, int]

IOW for generators the type written in the signature is what we get; but for native coroutines the type written is wrapped in Awaitable[] (which matches what await accepts). OTOH inside a generator, the type to be used in return statements is the third parameter of the Generator type written, while inside a native coroutine the return statement just matches the written return type.

This very long ramble is leading up to the question: What should @asyncio.coroutine to return?

I think we first need to define a further type: something that's acceptable by both await and yield from. Let's call this AwaitableGenerator[T] -- it is both Awaitable[T] and Generator[Any, None, T] (perhaps defined using multiple inheritance, perhaps using the Intersection primitive proposed in python/typing#213). The key property here is that it is acceptable to both await and yield from (and run_until_complete()).

So when @asyncio.coroutine wraps a native coroutine whose return type is written as T, and whose return type as seen by callers (when undecorated) is Awaitable[T], the decorator converts the return type to AwaitableGenerator[T] (but leaves the argument list unchanged -- this is the part we can't express yet using PEP 484).

And when it wraps a generator function whose return type is (written as, and seen by callers as) Generator[..., ..., T], the decorator converts it to AwaitableGenerator[T].

There's one more special case: when @asyncio.coroutine wraps a plain function returning T (where T is not Generator), the decorator turns it into a trivial (non-yielding) generator with type Generator[Any, None, T].

How to support asyncio.Future in await is relatively simple -- we can just make it inherit from Awaitable. (This is similar to how it's supported by yield from -- it inherits from Iterable and its __iter__ method is declared as Generator[Any, None, T].)

Finally a note about the difference between @types.coroutine and @asyncio.coroutine. The former does not have any notion of futures, and it doesn't special-case plain functions -- it requires a generator function, or a plain function that returns a generator object (this is a fairly obscure edge case that is nevertheless worth supporting, e.g. generators compiled with Cython).

As a stretch goal we might try to distinguish between the two decorators as follows: in asyncio, the first parameter to Generator must actually be asyncio.Future. Details are in Task._step in asyncio/tasks.py, around line 252.

[gvanrossum]

So, to be specific, if we had a feature to modify a Callable's return type, here's how I would define stubs for the two decorators.

Prerequisite: the AwaitableGenerator class:

T = TypeVar('T')
class AwaitableGenerator(Awaitable[T], Generator[Any, Any, T], Generic[T]):
    pass  # Nothing here

Then, @types.coroutine:

AA = TypeVar('AA')
T = TypeVar('T')
def coroutine(func: Callable[AA, Generator[Any, Any, T]]
             ) -> Callable[AA, AwaitableGenerator[T]: ...

Finally, @asyncio.coroutine, using overloading:

AA = TypeVar('AA')
T = TypeVar('T')
@overload
def coroutine(func: Callable[AA, Generator[Any, Any, T]]
             ) -> Callable[AA, AwaitableGenerator[T]]: ...
@overload
def coroutine(func: Callable[AA, Awaitable[T]]) -> Callable[AA, Awaitable[T]]: ...
@overload
def coroutine(func: Callable[AA, T]]) -> Callable[AA, Awaitable[T]]:

I suspect that the third variant of the overloading doesn't actually work, because it overlaps with the other two. More reasons to special-case this in mypy.

(To be clear, the notation using Callable[AA, ...] doesn't actually work -- it causes an error both in mypy and at runtime, since both mypy and the typing.Callable implementation insist that the first parameter is either an Ellipsis or a list of types.)

[gvanrossum]

A final question to myself: where should AwaitableGenerator be defined?

The obvious answer is to add it to the typing.pyi stub. Then we'll also have to add it to PEP 484 and to the typing.py implementation before people can define their own functions that return this type. But mypy itself won't have to wait for the latter changes, as long as we update typeshed. We'd still have to have agreement on PEP 484 though.

[gvanrossum]

I've got a bunch of this mostly working (https://github.com/python/mypy/tree/add-crawl2), but there's a remaining issue. I'm currently just updating the type of the decorated FuncItem to change its return type to AwaitableGenerator[...]. But that thwarts type-checking inside the function, e.g. like this:

@coroutine
def foo() -> Generator[int, None, str]:
    yield 'wrong!'
    return 'ok'

This should give an error on the yield, as it does without the @coroutine decorator -- but with it, the yield-type is replaced with Any (as is the receive-type).

Possible (hacky) solution: give AwaitableGenerator a second type parameter, and when updating the definition's type, set it to AwaitableGenerator[T, <original_type>]. Then get_generator_{yield,receive}_type() can extract the desired type from there.

[Update:] Or alternatively I could just use AwaitableGenerator[<original_type>], since this is really just a hack that stays inside mypy. (At least, for now that's what I'm thinking -- not to make this an actual definition in PEP 484.)

[gvanrossum]

FWIW I've got all this sorted out now in the add-crawl2 branch. No PR yet (I need to update PR #1808 in response to the latest code review first).

[gvanrossum]

Closing now that PR #1808 has been merged and the new PR #1946 (which implements the scheme I designed here) is ready for review.