Асинхронный декоратор для генераторов и сопрограмм

Question

Этот вопрос относится к тому же вопросу в синхронной версии .Цель состоит в том, чтобы разработать декоратор, который будет принимать в качестве входных данных либо генератор, либо сопрограмму в качестве параметра.Код, который у меня есть, выглядит следующим образом:

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)

Ошибка, которую это дает:

'return' with value in async generator

Я полагаю, что статически проверяется только тот факт, что decorated содержит оба yield и return со значением.

Есть ли способ заставить эту работу?(Кроме наличия в say_hello параметра, указывающего, является ли function генератором или сопрограммой).

Answer 1

Вы можете в основном использовать тот же механизм, который используется другим ответом, но применяется к сопрограммам.Например:

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

Обратите внимание, что в этом случае decorated определяется с использованием def вместо async def.Это гарантирует, что при вызове он сразу начинает работать и может выбирать, что возвращать, асинхронный генератор-итератор для объекта сопрограммы.Поскольку decorated возвращает объект, созданный путем вызова функции, определенной с async def inner, функционально он эквивалентен самому async def.

Ваша test функция неверна, поскольку онаиспользует for для перебора асинхронного генератора, и он перебирает сопрограмму.Вместо этого он должен быть async def и использовать async for для перебора асинхронного генератора и await для ожидания сопрограммы.Я использовал следующий код для проверки (generator и coroutine не изменились):

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