-
Notifications
You must be signed in to change notification settings - Fork 13
/
asyncache.py
127 lines (106 loc) · 4.36 KB
/
asyncache.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
"""
Helpers to use [cachetools](https://github.com/tkem/cachetools) with
asyncio.
"""
import functools
import inspect
__all__ = ["cached", "cachedmethod"]
import cachetools
# noinspection PyUnresolvedReferences
def cached(cache, key=cachetools.keys.hashkey, lock=None):
"""
Decorator to wrap a function or a coroutine with a memoizing callable
that saves results in a cache.
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 gets updated. If it wraps a coroutine, ``lock``
must implement ``__aenter__`` and ``__aexit__``.
"""
def decorator(func):
if inspect.iscoroutinefunction(func):
if cache is None:
async def wrapper(*args, **kwargs):
return await func(*args, **kwargs)
elif lock is None:
async def wrapper(*args, **kwargs):
k = key(*args, **kwargs)
try:
return cache[k]
except KeyError:
pass # key not found
v = await func(*args, **kwargs)
try:
cache[k] = v
except ValueError:
pass # value too large
return v
else:
async def wrapper(*args, **kwargs):
k = key(*args, **kwargs)
try:
async with lock:
return cache[k]
except KeyError:
pass # key not found
v = await func(*args, **kwargs)
# in case of a race, prefer the item already in the cache
try:
async with lock:
return cache.setdefault(k, v)
except ValueError:
return v # value too large
return functools.update_wrapper(wrapper, func)
else:
raise NotImplementedError("Use cachetools.cached for non-async functions")
return decorator
# noinspection PyUnresolvedReferences
def cachedmethod(cache, key=cachetools.keys.hashkey, lock=None):
"""Decorator to wrap a class or instance method with a memoizing
callable that saves results in a cache.
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 gets updated. If it wraps a coroutine, ``lock``
must implement ``__aenter__`` and ``__aexit__``.
"""
def decorator(method):
if inspect.iscoroutinefunction(method):
if lock is None:
async def wrapper(self, *args, **kwargs):
c = cache(self)
if c is None:
return await method(self, *args, **kwargs)
k = key(*args, **kwargs)
try:
return c[k]
except KeyError:
pass # key not found
v = await method(self, *args, **kwargs)
try:
c[k] = v
except ValueError:
pass # value too large
return v
else:
async def wrapper(self, *args, **kwargs):
c = cache(self)
if c is None:
return await method(self, *args, **kwargs)
k = key(*args, **kwargs)
try:
async with lock(self):
return c[k]
except KeyError:
pass # key not found
v = await method(self, *args, **kwargs)
# in case of a race, prefer the item already in the cache
try:
async with lock(self):
return c.setdefault(k, v)
except ValueError:
return v # value too large
return functools.update_wrapper(wrapper, method)
else:
raise NotImplementedError(
"Use cachetools.cachedmethod for non-async methods"
)
return decorator