Search code examples
pythonasynchronouscachingaiohttp

Ignoring a parameter while caching asyncio functions


So in aiohttp it is suggested to reuse the ClientSession to avoid overhead.

So say I have something like this:

async def get_id_from_url(session, url):
    async with session.get(url) as response:
        return (await response.json())['id'] if response.ok else ''


async def get_ids(urls: List[str]):
    async with aiohttp.ClientSession() as session:
        ids = await asyncio.gather(*[get_id_from_url(session, url) for url in urls])
        print(ids)

I would like to cache the get_id_from_url function, I have looked at many questions/packages relating to caching async functions: aiocache and this answer.

But there is one problem, I don't want the caching to depend on the session parameter, I only want it to depend on the url parameter.

How can I do that?


Solution

  • After doing a lot of research I found the cachetools library and this pull request.

    I just took the solution from there: this is an implementation of asyncio support for cachetools. And cachetools already allows you to pass a key parameter that allows you to chose which parameters of the function will be involved in the hashing.

    import functools
    import inspect
    import logging
    
    from cachetools.keys import hashkey
    
    logger = logging.getLogger(__name__)
    
    
    class NullContext(object):
        """A class for noop context managers."""
    
        def __enter__(self):
            """Return ``self`` upon entering the runtime context."""
            return self
    
        def __exit__(self, exc_type, exc_value, traceback):
            """Raise any exception triggered within the runtime context."""
    
        async def __aenter__(self):
            """Return ``self`` upon entering the runtime context."""
            return self
    
        async def __aexit__(self, exc_type, exc_value, traceback):
            """Raise any exception triggered within the runtime context."""
    
    
    def aiocached(cache, key=hashkey, lock=None):
        """Decorator to wrap a function or a coroutine with a memoizing callable.
    
        When ``lock`` is provided for a standard function, it's expected to
        implement ``__enter__`` and ``__exit__`` that will be used to lock
        the cache when it gets updated. If it wraps a coroutine, ``lock``
        must implement ``__aenter__`` and ``__aexit__``.
        """
        lock = lock or NullContext()
    
        def decorator(func):
            if not inspect.iscoroutinefunction(func):
                raise RuntimeError('Use aiocached only with async functions')
    
            async def wrapper(*args, **kwargs):
                fk = key(*args, **kwargs)
                async with lock:
                    fval = cache.get(fk)
                # cache hit
                if fval is not None:
                    return fval
                # cache miss
                fval = await func(*args, **kwargs)
                try:
                    async with lock:
                        cache[fk] = fval
                except ValueError:
                    logger.debug('Failed to cache {0}'.format(fval))
                return fval
            return functools.wraps(func)(wrapper)
        return decorator
    

    Now you can use the aiocached decorator just like cached