如何为异步协程编写高阶函数augment_async？附扩展问题

先帮你拆解这个问题，分两步搞定：先实现支持异步函数的基础增强功能，再扩展支持直接增强已调用的协程对象。

1. 基础版：增强异步函数的实现 #

原augment函数失效的核心原因是：异步函数调用后返回的是协程对象，不是直接的字典结果，直接用**func(x)解包会触发TypeError，同时未被await的协程会抛出RuntimeWarning。

我们只需要把增强后的函数也改成异步的，先等待原异步函数执行完成拿到结果，再合并metadata：

async def augment_async(func):
    async def augmented_function(x):
        # 先await异步函数的结果，拿到真正的字典
        original_result = await func(x)
        return {**original_result, "metadata": "test"}
    return augmented_function

# 测试代码
async def f(x):
    return {"x": x}

# 使用方式
coroutine = augment_async(f)(1)
final_result = await coroutine
print(final_result)  # 输出: {"x": 1, "metadata": "test"}

这样调用augment_async(f)(1)会返回一个协程对象，await之后就能得到包含metadata的完整字典，完全符合你的基础需求。

2. 扩展版：支持增强已调用的协程对象 #

接下来实现扩展需求：让await augment_async(f(1))也能返回正确结果。这时候需要让augment_async能自动识别输入是异步函数还是已调用后的协程对象，分别处理：

import inspect

async def augment_async(obj):
    # 情况1：传入的是异步函数（async def定义的可调用对象）
    if inspect.iscoroutinefunction(obj):
        async def augmented_function(x):
            original_result = await obj(x)
            return {**original_result, "metadata": "test"}
        return augmented_function
    # 情况2：传入的是已经调用后的协程对象
    elif inspect.iscoroutine(obj):
        original_result = await obj
        return {**original_result, "metadata": "test"}
    # 其他情况抛出明确错误
    else:
        raise TypeError("仅支持异步函数或协程对象作为输入")

# 两种场景测试
async def f(x):
    return {"x": x}

# 场景1：增强函数后调用
result1 = await augment_async(f)(1)
print(result1)  # {"x": 1, "metadata": "test"}

# 场景2：直接增强已调用的协程
result2 = await augment_async(f(1))
print(result2)  # {"x": 1, "metadata": "test"}

这个版本的augment_async会自动判断输入类型：如果是异步函数，就返回增强后的异步函数；如果是协程对象，就直接执行它并添加metadata返回，完美覆盖两种使用场景。

内容的提问来源于stack exchange，提问作者jhoareau