cache.py

import trio
import weakref
import cachetools

# Note: this assumes you only have one trio thread running at once...
# In the unlikely event that you have multiple, you should put the locks and cache
# into thread-local or run-local storage.
def trio_lru_cache(fn):
    cache = cachetools.LRUCache()
    # Trick: this contains a lock for each unique cache key in use,
    # that we use to make sure we don't have multiple tasks trying
    # to rebuild the same cache key at once. Since it's a WeakValueDictionary,
    # unused keys are automatically evicted by the GC.
    locks = weakref.WeakValueDictionary()
  
    @functools.wraps(fn)
    async def wrapper(*args, **kwargs):
        key = make_key(args, kwargs)
        try:
            return cache[key]
        except KeyError:
            lock = locks.setdefault(key, trio.Lock())
            async with lock:
                try:
                    return cache[key]
                except KeyError:
                    new_value = await fn(*args, **kwargs)
                    cache[key] = new_value
                    return new_value
    
    return wrapper