Asynchronous decorator for both generators and coroutines

Question

This question related to the same question in a synchronous version. The goal is to design decorator that would take as input either a generator or a coroutine as parameter. The code I have looks like:

import asyncio


def say_hello(function):

    async def decorated(*args, **kwargs):
        return_value = function(*args, **kwargs)
        if isinstance(return_value, types.AsyncGeneratorType):
            print("Hello async generator!")
            async for v in return_value:
                yield v
        else:
            print("Hello coroutine!")
            return await return_value

    return decorated


@helpers.say_hello
async def generator():
    for i in range(5):
        await asyncio.sleep(0.2)
        yield i

@helpers.say_hello
async def coroutine():
    await asyncio.sleep(1)
    return list(range(5))


async def test():
    for v in generator():
        print(v)
    for v in coroutine():
        print(v)

The error this gives is:

'return' with value in async generator

Which I guess is just checked statically on the fact that decorated contains both a yield and a return with a value.

Is there any way to make this work? (Other than having a parameter in say_hello to specify if function is a generator or a coroutine).

Answer 1

You can basically use the same mechanism deployed by the other answer, but applied for coroutines. For example:

def say_hello(function):
    def decorated(*args, **kwargs):
        function_instance = function(*args, **kwargs)
        if isinstance(function_instance, types.AsyncGeneratorType):
            async def inner():
                print("Hello async generator!")
                async for v in function_instance:
                    yield v
        else:
            async def inner():
                print("Hello coroutine!")
                return await function_instance
        return inner()
    return decorated

Note that in this case decorated is defined using def rather than async def. This ensures that, when called, it immediately starts running and is able to choose what to return, an async generator-iterator on a coroutine object. Since decorated returns an object created by calling a function defined with async def inner, it is functionally equivalent to being async def itself.

Your test function is incorrect because it uses for to iterate over an async generator, and it iterates over the coroutine. Instead, it should be async def and use async for to iterate over the async generator and await to await the coroutine. I used the following code to test (generator and coroutine are unchanged):

async def test():
    async for v in generator():
        print(v)
    await coroutine()

asyncio.run(test())
# or, on Python 3.6 and older:
#asyncio.get_event_loop().run_until_complete(test())