Wrapping an asynchronous generator in Python

Question

Since Python 3.6 and PEP 525 one can use asynchronous generator:

import asyncio
async def asyncgen():
 yield 1
 yield 2
async def main():
 async for i in asyncgen():
 print(i)
asyncio.run(main())

I created a function which is able to wrap any asynchronous generator, the same way you would wrap a basic function using @decorator.

def asyncgen_wrapper(generator):
 async def wrapped(*args, **kwargs):
 print("The wrapped asynchronous generator is iterated")
 gen = generator(*args, **kwargs)
 try:
 value = await gen.__anext__()
 except StopAsyncIteration:
 return
 while True:
 to_send = yield value
 try:
 value = await gen.asend(to_send)
 except StopAsyncIteration:
 return
 return wrapped

Wrapping an asynchronous generator seems quite complicated compared to wrapping a basic generator:

def gen_wrapper(generator):
 def wrapped(*args, **kwargs):
 return (yield from generator(*args, **kwargs))
 return wrapped

I mainly concerned about correctness of my wrapper. I want the wrapper to be as transparent as possible regarding the wrapped generator. Which leads me to have two questions:

• Is there a more straightforward way to implement a decorator for asynchronous generators?
• Does my implementation handle all possible edge cases (think to asend() for example)?

Answer 1

"Correct" is nebulous. If it were to mean here something like "the semantics of any wrapped generator through the public API is the same as if it were not wrapped, apart from a print statement before the generator is created", then, no, the code isn't correct since it doesn't preserve the effects of raising an exception or closing the generator. For example,

import asyncio
class StopItNow( Exception ): 
 pass
async def f():
 try:
 yield 0
 except StopItNow:
 yield 1
async def g( f ): 
 x = f()
 print( await x.asend( None ) )
 print( await x.athrow( StopItNow ) )
asyncio.run( g( f ) )
asyncio.run( g( asyncgen_wrapper( f ) ) )

. I think it would be a non-trivial undertaking to "correctly" implement a wrapper. Interestingly, in the same PEP (525) that you've linked in your question,

While it is theoretically possible to implement yield from support for asynchronous generators, it would require a serious redesign of the generators implementation.

https://www.python.org/dev/peps/pep-0525/#asynchronous-yield-from

. You might find that it's a whole lot easier to implement some sort of injection or parameterise the generator itself to accept callables.

Otherwise, from my experimentation and taking hints from PEP 380, there are implementation details that are omitted from PEP 525 that would be necessary to emulate the behaviour of an unwrapped generator.

This is the result of some toying around:

import functools
import sys
def updated_asyncgen_wrapper( generator ): 
 @functools.wraps( generator )
 async def wrapped( *args, **kwargs ): 
 print( "iterating wrapped generator" )
 gen = generator( *args, **kwargs )
 to_send, is_exc = ( None, ), False
 while True:
 try:
 do = gen.athrow if is_exc else gen.asend
 value = await do( *to_send )
 except StopAsyncIteration:
 return
 try:
 to_send, is_exc = ( (yield value), ), False
 except GeneratorExit:
 await gen.aclose()
 raise
 except:
 to_send, is_exc = sys.exc_info(), True
 return wrapped

. This isn't "correct" either, since it doesn't disambiguate between an attempt to close the generator and an explicit throw of an instance of GeneratorExit, which is, though, marked distinctly for usage by the former case. This might be good enough for internal use.