如何让检测慢请求的aiohttp装饰器支持异步上下文管理器？

问题场景 #

已实现report_slow_call装饰器用于检测aiohttp慢请求，通过猴子补丁修改ClientSession.get方法后，直接await调用正常，但使用async with异步上下文管理器调用时触发错误：

RuntimeWarning: coroutine 'report_slow_call.<locals>.decorator.<locals>.wrapper' was never awaited
TypeError: 'coroutine' object does not support the asynchronous context manager protocol

错误原因 #

原ClientSession.get方法返回的ClientResponse对象同时具备两个特性：

1. 实现__await__方法，支持await调用；
2. 实现异步上下文管理器协议（__aenter__/__aexit__），支持async with语法。

而原装饰器将get方法包装为异步函数，调用后返回的是协程对象，协程对象仅支持await，无法直接作为异步上下文管理器使用（低版本Python不支持async with直接处理协程）。

解决方案 #

修改装饰器，让包装后的get方法返回一个同时支持await和异步上下文管理器协议的代理对象，完全兼容原方法的两种调用方式。

修正后的装饰器代码 #

import time
from functools import wraps

def report_slow_call(seconds=1.0):
    def decorator(func):
        class WrappedResponse:
            def __init__(self, timed_coro):
                self.timed_coro = timed_coro
                self.response = None

            async def __await__(self):
                self.response = await self.timed_coro
                return self.response

            async def __aenter__(self):
                self.response = await self.timed_coro
                return await self.response.__aenter__()

            async def __aexit__(self, exc_type, exc_val, exc_tb):
                return await self.response.__aexit__(exc_type, exc_val, exc_tb)

        @wraps(func)
        def wrapper(*args, **kwargs):
            start_ts = time.perf_counter()
            # 包装原协程，添加计时逻辑
            async def timed_coro():
                nonlocal start_ts
                response = await func(*args, **kwargs)
                elapsed = time.perf_counter() - start_ts
                # 提取请求URL
                url = args[0] if args else kwargs.get('url')
                if elapsed > seconds:
                    print(f"Slow call to: {url} ({elapsed:.2f} seconds).")
                return response
            # 返回代理对象，兼容两种调用方式
            return WrappedResponse(timed_coro())
        return wrapper
    return decorator

使用示例 #

import aiohttp

# 初始化ClientSession并应用装饰器
http_client = aiohttp.ClientSession()
http_client.get = report_slow_call(seconds=0.03)(http_client.get)

# 方式1：直接await调用
await http_client.get("http://stackoverflow.com")

# 方式2：async with上下文管理器调用
async with http_client.get("http://stackoverflow.com") as response:
    print(f"Response status: {response.status}")
    response.raise_for_status()

说明 #

• WrappedResponse类作为代理，同时实现__await__、__aenter__、__aexit__方法，完全复刻原ClientResponse的行为；
• 计时逻辑封装在timed_coro中，不影响原请求的正常流程；
• 兼容Python 3.6及以上版本，同时支持两种调用方式。

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