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

Question

此问题与synchronous version 中的同一问题相关。 目标是设计将生成器或协程作为参数作为输入的装饰器。我的代码如下：

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