如何正确包装协程的 __await__ 方法以实现异步资源操作的日志追踪

在Python异步编程中，协程的异步执行特性使得传统的同步日志追踪方式难以覆盖资源操作的完整生命周期，而直接修改协程内部逻辑会增加代码耦合度。包装协程的__await__方法是一种低侵入性的解决方案，可以在不改动原有业务代码的前提下，实现异步资源操作的日志自动记录。

__await__方法的基本工作原理

Python中的协程对象实现了__await__方法，该方法返回一个迭代器，用于驱动协程的执行流程。当使用await关键字等待协程时，解释器会自动调用协程对象的__await__方法，逐步执行协程内部的逻辑。我们可以通过自定义__await__方法的返回值，在原有协程执行前后插入额外的逻辑，比如日志记录。

包装__await__方法的实现思路

包装__await__方法的核心思路是创建一个包装类，该类的实例可以替代原有协程对象，同时重写__await__方法，在调用原有协程的__await__方法前后添加日志逻辑。具体步骤如下：

• 创建包装类，接收原始协程对象作为参数
• 重写包装类的__await__方法，先记录操作开始日志
• 调用原始协程的__await__方法获取迭代器
• 迭代执行原始协程逻辑，捕获可能的异常并记录
• 操作完成后记录结束日志，返回最终结果

基础实现代码示例

以下是一个通用的协程包装实现，支持对异步资源操作的开始、结束、异常情况进行日志追踪：

import logging
import time

# 配置基础日志
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

class AwaitableWrapper:
    def __init__(self, coro, resource_name):
        self.coro = coro  # 原始协程对象
        self.resource_name = resource_name  # 资源名称，用于日志标识

    def __await__(self):
        # 记录操作开始日志
        start_time = time.time()
        logger.info(f"异步资源操作开始，资源名称：{self.resource_name}")
        try:
            # 获取原始协程的迭代器
            coro_iter = self.coro.__await__()
            # 迭代执行原始协程逻辑
            result = yield from coro_iter
            # 记录操作成功结束日志
            cost_time = (time.time() - start_time) * 1000
            logger.info(f"异步资源操作成功结束，资源名称：{self.resource_name}，耗时：{cost_time:.2f}ms")
            return result
        except Exception as e:
            # 记录操作异常日志
            logger.error(f"异步资源操作异常，资源名称：{self.resource_name}，异常信息：{str(e)}")
            raise  # 重新抛出异常，不影响原有逻辑

# 装饰器函数，方便快速包装协程
def trace_async_resource(resource_name):
    def decorator(func):
        def wrapper(*args, **kwargs):
            coro = func(*args, **kwargs)
            return AwaitableWrapper(coro, resource_name)
        return wrapper
    return decorator

实际应用场景示例

假设我们有一个异步读取文件资源的操作，使用上述包装方案实现日志追踪：

import asyncio

# 使用装饰器包装异步资源操作函数
@trace_async_resource("本地文件读取")
async def read_async_file(file_path):
    await asyncio.sleep(0.5)  # 模拟异步读取耗时
    # 模拟读取逻辑，这里省略实际文件读取代码
    return f"文件{file_path}的内容"

async def main():
    try:
        content = await read_async_file("/data/test.txt")
        print(f"读取结果：{content}")
    except Exception as e:
        print(f"操作失败：{e}")

if __name__ == "__main__":
    asyncio.run(main())

运行上述代码后，会输出类似以下的日志：

2024-05-20 14:30:00,123 - INFO - 异步资源操作开始，资源名称：本地文件读取
2024-05-20 14:30:00,624 - INFO - 异步资源操作成功结束，资源名称：本地文件读取，耗时：501.23ms
读取结果：文件/data/test.txt的内容

注意事项

在实际使用包装__await__方法实现日志追踪时，需要注意以下几点：

• 包装类必须正确实现__await__方法，否则无法被await关键字识别
• 异常需要重新抛出，避免影响原有协程的错误处理逻辑
• 如果原始协程有其他特殊属性或方法，包装类需要做对应的转发处理，避免属性丢失
• 日志内容可以根据实际需求扩展，比如添加操作参数、调用栈信息等

适配不同异步框架的优化

如果使用的是特定异步框架（如aiohttp、asyncpg等），可以针对框架的协程特性做进一步优化。例如针对aiohttp的异步请求操作，可以在日志中添加请求URL、响应状态码等信息：

@trace_async_resource("aiohttp异步请求")
async def fetch_url(session, url):
    async with session.get(url) as response:
        # 可以在包装逻辑中扩展获取响应信息，这里通过协程内部返回额外数据的方式传递
        content = await response.text()
        return {
            "content": content,
            "status": response.status,
            "url": url
        }

# 优化后的AwaitableWrapper可以在日志中记录响应状态码
class EnhancedAwaitableWrapper(AwaitableWrapper):
    def __await__(self):
        start_time = time.time()
        logger.info(f"异步资源操作开始，资源名称：{self.resource_name}")
        try:
            coro_iter = self.coro.__await__()
            result = yield from coro_iter
            cost_time = (time.time() - start_time) * 1000
            # 如果结果包含状态码，添加到日志中
            if isinstance(result, dict) and "status" in result:
                logger.info(f"异步资源操作成功结束，资源名称：{self.resource_name}，耗时：{cost_time:.2f}ms，响应状态码：{result['status']}")
            else:
                logger.info(f"异步资源操作成功结束，资源名称：{self.resource_name}，耗时：{cost_time:.2f}ms")
            return result
        except Exception as e:
            logger.error(f"异步资源操作异常，资源名称：{self.resource_name}，异常信息：{str(e)}")
            raise

协程__await__方法异步资源操作日志追踪Python异步编程修改时间：2026-07-04 00:49:03