192 lines
4.6 KiB
Python
192 lines
4.6 KiB
Python
import asyncio
|
|
from collections import OrderedDict
|
|
from functools import _CacheInfo, _make_key, partial, wraps
|
|
|
|
__version__ = "1.0.3"
|
|
|
|
__all__ = ("alru_cache",)
|
|
|
|
|
|
def unpartial(fn):
|
|
while hasattr(fn, "func"):
|
|
fn = fn.func
|
|
|
|
return fn
|
|
|
|
|
|
def _done_callback(fut, task):
|
|
if task.cancelled():
|
|
fut.cancel()
|
|
return
|
|
|
|
exc = task.exception()
|
|
if exc is not None:
|
|
fut.set_exception(exc)
|
|
return
|
|
|
|
fut.set_result(task.result())
|
|
|
|
|
|
def _cache_invalidate(wrapped, typed, *args, **kwargs):
|
|
key = _make_key(args, kwargs, typed)
|
|
|
|
exists = key in wrapped._cache
|
|
|
|
if exists:
|
|
wrapped._cache.pop(key)
|
|
|
|
return exists
|
|
|
|
|
|
def _cache_clear(wrapped):
|
|
wrapped.hits = wrapped.misses = 0
|
|
wrapped._cache = OrderedDict()
|
|
wrapped.tasks = set()
|
|
|
|
|
|
def _open(wrapped):
|
|
if not wrapped.closed:
|
|
raise RuntimeError("alru_cache is not closed")
|
|
|
|
was_closed = (
|
|
wrapped.hits == wrapped.misses == len(wrapped.tasks) == len(wrapped._cache) == 0
|
|
)
|
|
|
|
if not was_closed:
|
|
raise RuntimeError("alru_cache was not closed correctly")
|
|
|
|
wrapped.closed = False
|
|
|
|
|
|
def _close(wrapped, *, cancel=False, return_exceptions=True):
|
|
if wrapped.closed:
|
|
raise RuntimeError("alru_cache is closed")
|
|
|
|
wrapped.closed = True
|
|
|
|
if cancel:
|
|
for task in wrapped.tasks:
|
|
if not task.done(): # not sure is it possible
|
|
task.cancel()
|
|
|
|
return _wait_closed(wrapped, return_exceptions=return_exceptions)
|
|
|
|
|
|
async def _wait_closed(wrapped, *, return_exceptions):
|
|
wait_closed = asyncio.gather(*wrapped.tasks, return_exceptions=return_exceptions)
|
|
|
|
wait_closed.add_done_callback(partial(_close_waited, wrapped))
|
|
|
|
ret = await wait_closed
|
|
|
|
# hack to get _close_waited callback to be executed
|
|
await asyncio.sleep(0)
|
|
|
|
return ret
|
|
|
|
|
|
def _close_waited(wrapped, _):
|
|
wrapped.cache_clear()
|
|
|
|
|
|
def _cache_info(wrapped, maxsize):
|
|
return _CacheInfo(
|
|
wrapped.hits,
|
|
wrapped.misses,
|
|
maxsize,
|
|
len(wrapped._cache),
|
|
)
|
|
|
|
|
|
def __cache_touch(wrapped, key):
|
|
try:
|
|
wrapped._cache.move_to_end(key)
|
|
except KeyError: # not sure is it possible
|
|
pass
|
|
|
|
|
|
def _cache_hit(wrapped, key):
|
|
wrapped.hits += 1
|
|
__cache_touch(wrapped, key)
|
|
|
|
|
|
def _cache_miss(wrapped, key):
|
|
wrapped.misses += 1
|
|
__cache_touch(wrapped, key)
|
|
|
|
|
|
def alru_cache(
|
|
fn=None,
|
|
maxsize=128,
|
|
typed=False,
|
|
*,
|
|
cache_exceptions=True,
|
|
):
|
|
def wrapper(fn):
|
|
_origin = unpartial(fn)
|
|
|
|
if not asyncio.iscoroutinefunction(_origin):
|
|
raise RuntimeError("Coroutine function is required, got {}".format(fn))
|
|
|
|
# functools.partialmethod support
|
|
if hasattr(fn, "_make_unbound_method"):
|
|
fn = fn._make_unbound_method()
|
|
|
|
@wraps(fn)
|
|
async def wrapped(*fn_args, **fn_kwargs):
|
|
if wrapped.closed:
|
|
raise RuntimeError("alru_cache is closed for {}".format(wrapped))
|
|
|
|
loop = asyncio.get_event_loop()
|
|
|
|
key = _make_key(fn_args, fn_kwargs, typed)
|
|
|
|
fut = wrapped._cache.get(key)
|
|
|
|
if fut is not None:
|
|
if not fut.done():
|
|
_cache_hit(wrapped, key)
|
|
return await asyncio.shield(fut)
|
|
|
|
exc = fut._exception
|
|
|
|
if exc is None or cache_exceptions:
|
|
_cache_hit(wrapped, key)
|
|
return fut.result()
|
|
|
|
# exception here and cache_exceptions == False
|
|
wrapped._cache.pop(key)
|
|
|
|
fut = loop.create_future()
|
|
task = loop.create_task(fn(*fn_args, **fn_kwargs))
|
|
task.add_done_callback(partial(_done_callback, fut))
|
|
|
|
wrapped.tasks.add(task)
|
|
task.add_done_callback(wrapped.tasks.remove)
|
|
|
|
wrapped._cache[key] = fut
|
|
|
|
if maxsize is not None and len(wrapped._cache) > maxsize:
|
|
wrapped._cache.popitem(last=False)
|
|
|
|
_cache_miss(wrapped, key)
|
|
return await asyncio.shield(fut)
|
|
|
|
_cache_clear(wrapped)
|
|
wrapped._origin = _origin
|
|
wrapped.closed = False
|
|
wrapped.cache_info = partial(_cache_info, wrapped, maxsize)
|
|
wrapped.cache_clear = partial(_cache_clear, wrapped)
|
|
wrapped.invalidate = partial(_cache_invalidate, wrapped, typed)
|
|
wrapped.close = partial(_close, wrapped)
|
|
wrapped.open = partial(_open, wrapped)
|
|
|
|
return wrapped
|
|
|
|
if fn is None:
|
|
return wrapper
|
|
|
|
if callable(fn) or hasattr(fn, "_make_unbound_method"):
|
|
return wrapper(fn)
|
|
|
|
raise NotImplementedError("{} decorating is not supported".format(fn))
|