Python decorator for aiohttp to work as the asynchronius context manager

Question

I am trying to write a wrapper for aiohttp to detect slow queries.

Here is my decorator:

import time

def report_slow_call(seconds=1.0):
    def decorator(func):
        async def wrapper(*args, **kwargs):
            start_ts = time.perf_counter()
            result = await func(*args, **kwargs)
            elapsed = time.perf_counter() - start_ts
            url = args[0] if args else kwargs.get('url')
            if elapsed > seconds:
                print(f"Slow call to: {url} ({elapsed:.2f} seconds).")
            return result
        return wrapper
    return decorator

I monkey patch client object get function, and it's working fine (executed in ipython):

In [2]: import aiohttp
   ...: http_client = aiohttp.ClientSession()
   ...: http_client.get = report_slow_call(seconds=0.03)(http_client.get)
   ...: await http_client.get("http://stackoveflow.com")
Slow call to: http://stackoveflow.com (0.11 seconds).
Out[2]: 
<ClientResponse(http://stackoveflow.com) [200 OK]>
<CIMultiDictProxy('accept-ch': 'Sec-CH-UA, Sec-CH-UA-Platform, Sec-CH-UA-Platform-Version, Sec-CH-UA-Mobile', 'Cache-Control': 'max-age=0, private, must-revalidate', 'Connection': 'close', 'Content-Length': '477', 'Content-Type': 'text/html; charset=utf-8', 'Date': 'Mon, 17 Apr 2023 15:42:50 GMT', 'Server': 'nginx', 'Set-Cookie': 'sid=82900a62-dd36-11ed-85ae-a4d3e99a9525; path=/; domain=.stackoveflow.com; expires=Sat, 05 May 2091 18:56:57 GMT; max-age=2147483647; HttpOnly')>

Problem: When I am trying to run it through the context manager, I am getting the following error:

In [4]: async with http_client.get("http://stackoveflow.com", raise_for_status=True) as response:
   ...:     print(response)
   ...: 
<ipython-input-4-fd5ea6565740>:1: RuntimeWarning: coroutine 'report_slow_call.<locals>.decorator.<locals>.wrapper' was never awaited
  async with http_client.get("http://stackoveflow.com", raise_for_status=True) as response:
RuntimeWarning: Enable tracemalloc to get the object allocation traceback
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[4], line 1
----> 1 async with http_client.get("http://stackoveflow.com", raise_for_status=True) as response:
      2     print(response)

TypeError: 'coroutine' object does not support the asynchronous context manager protocol

Would anyone of you suggest an idea for proper implementation, please?

Answer 1

I checked sources of httpaio and learnt how functions are executed.

    def get(
        self, url: StrOrURL, *, allow_redirects: bool = True, **kwargs: Any
    ) -> "_RequestContextManager":
        """Perform HTTP GET request."""
        return _RequestContextManager(
            self._request(hdrs.METH_GET, url, allow_redirects=allow_redirects, **kwargs)
        )

So, instead of patching ClientSession.get I patched ClientSession._request and decorator is working fine now.

Solution

import time
import aiohttp

def report_slow_call(seconds=1.0):
    def decorator(func):
        async def wrapper(*args, **kwargs):
            start_ts = time.perf_counter()
            result = await func(*args, **kwargs)
            elapsed = time.perf_counter() - start_ts
            url = args[0] if args else kwargs.get('url')
            if elapsed > seconds:
                print(f"Slow call to: {url} ({elapsed:.2f} seconds.)")
            return result
        return wrapper
    return decorator

url = "https://duckduckgo.com"
http_client = aiohttp.ClientSession()

# Monkey patch object's _request method
http_client._request = report_slow_call(seconds=0.01)(http_client._request)

# Execute ClientSession's functions
await http_client.get(url)

async with http_client.get(url, raise_for_status=True) as response:
    print(response)