Python async-generator not async
The point of async
/await
is to interleave tasks, not functions/generators. For example, when you await asyncio.sleep(1)
, your current coroutine is delayed along with the sleep. Similarly, an async for
delays its coroutine until the next item is ready.
In order to run your separate functionality, you must create each part as a separate task. Use a Queue
to exchange items between them - tasks will only be delayed until they have exchanged an item.
from asyncio import Queue, sleep, run, gather
# the original async generator
async def g():
for i in range(3):
await sleep(1)
yield i
async def producer(queue: Queue):
async for i in g():
print('send', i)
await queue.put(i) # resume once item is fetched
await queue.put(None)
async def consumer(queue: Queue):
x = await queue.get() # resume once item is fetched
while x is not None:
print('got', x)
await sleep(2)
x = await queue.get()
async def main():
queue = Queue()
# tasks only share the queue
await gather(
producer(queue),
consumer(queue),
)
run(main())
If you regularly need this functionality, you can also put it into a helper object that wraps an asynchronous iterable. The helper encapsulates the queue and separate task. You can apply the helper directly on an async iterable in an async for
statement.
from asyncio import Queue, sleep, run, ensure_future
# helper to consume iterable as concurrent task
async def _enqueue_items(async_iterable, queue: Queue, sentinel):
async for item in async_iterable:
await queue.put(item)
await queue.put(sentinel)
async def concurrent(async_iterable):
"""Concurrently fetch items from ``async_iterable``"""
queue = Queue()
sentinel = object()
consumer = ensure_future( # concurrently fetch items for the iterable
_enqueue_items(async_iterable, queue, sentinel)
)
try:
item = await queue.get()
while item is not sentinel:
yield item
item = await queue.get()
finally:
consumer.cancel()
# the original generator
async def g():
for i in range(3):
await sleep(1)
yield i
# the original main - modified with `concurrent`
async def main():
async for x in concurrent(g()):
print(x)
await sleep(2)
run(main())
As an alternative to doing this with a Queue, instead this solution chain Futures together, so that a Future's result is the current item and another Future to retrieve the next item (sort of like a linked list, so to speak):
from asyncio import sleep, get_event_loop, run, create_task
async def aiter(fut, async_generator):
try:
async for item in async_generator:
fut, prev_fut = get_event_loop().create_future(), fut
prev_fut.set_result((item, fut))
else:
fut.set_exception(StopAsyncIteration())
except Exception as e:
fut.set_exception(e)
async def concurrent(async_generator):
fut = get_event_loop().create_future()
create_task(aiter(fut, async_generator))
try:
while True:
item, fut = await fut
yield item
except StopAsyncIteration as e:
return
As an added bonus this solution will correctly handle exception that happens in g() by reraising the exception in the main() method with a traceback that will be useful for debugging.