用于生成器和协程的异步装饰器

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我猜这只是静态检查包含 ayield和 areturn的值的事实。

有什么办法可以让这个工作吗？（除了有一个参数来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())