Search code examples
pythonpython-asynciogenerator

Python async coroutine and generator in one object


Is there any possibility to have two coroutines, one iterating over an async generator, the other one awaiting the Stop of the Iteration? like

coro1(gen):
    async for _ in gen:
        await asyncio.sleep(1)

coro2(gen):
    await gen
    print('generation finished')

Solution

  • There's no automatic notification for that, but you can have coro1 notify coro2 with an asyncio.Event:

    import asyncio
    
    async def coro1(gen, event):
        async for _ in gen:
            await asyncio.sleep(1)
        event.set()
    
    async def coro2(event):
        await event.wait()
        print('generation finished')
    
    async def generator():
        yield 1
        yield 2
    
    async def main():
        event = asyncio.Event()
        gen = generator()
        await asyncio.gather(coro1(gen, event), coro2(event))
    
    asyncio.run(main())
    

    You can also build this notification into the generator itself, or put it into a wrapper generator:

    import asyncio
    import functools
    
    async def _notification_wrapper(gen, event):
        # Unfortunately, async generators don't support yield from.
        try:
            async for item in gen:
                yield item
        finally:
            event.set()
    
    def notification_wrapper(gen):
        event = asyncio.Event()
        return _notification_wrapper(gen, event), event
    
    async def generator():
        yield 1
        yield 2
    
    async def coro1(gen):
        async for _ in gen:
            await asyncio.sleep(1)
    
    async def coro2(event):
        await event.wait()
        print('generation finished')
    
    async def main():
        gen, event = notification_wrapper(generator())
    
        await asyncio.gather(coro1(gen), coro2(event))
    
    asyncio.run(main())
    

    Note that a generator (async or otherwise) cannot detect when code breaks out of a loop over the generator. If you want coro2 to stop waiting when coro1 breaks, you will have to have coro1 handle the notification, rather than trying to put it in the generator.