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)?

two_pc_turkey_breast_mince · Accepted Answer · 2020-05-07 21:38:08Z

"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.

Stack Exchange Network

Wrapping an asynchronous generator in Python

1 Answer 1

You must log in to answer this question.

Hot Network Questions

Wrapping an asynchronous generator in Python

1 Answer 1

You must log in to answer this question.

Related

Hot Network Questions