1from functools import wraps
2import time
3try:
4 from collections.abc import MutableMapping # Python 3.3+
5except ImportError:
6 from collections import MutableMapping # Python 2.7+
7import heapq
8from threading import Lock
9
10
11class _ExpiringMapping(MutableMapping):
12 _INDEX = "_index_"
13
14 def __init__(self, mapping=None, capacity=None, expires_in=None, lock=None,
15 *args, **kwargs):
16 """Items in this mapping can have individual shelf life,
17 just like food items in your refrigerator have their different shelf life
18 determined by each food, not by the refrigerator.
19
20 Expired items will be automatically evicted.
21 The clean-up will be done at each time when adding a new item,
22 or when looping or counting the entire mapping.
23 (This is better than being done indecisively by a background thread,
24 which might not always happen before your accessing the mapping.)
25
26 This implementation uses no dependency other than Python standard library.
27
28 :param MutableMapping mapping:
29 A dict-like key-value mapping, which needs to support __setitem__(),
30 __getitem__(), __delitem__(), get(), pop().
31
32 The default mapping is an in-memory dict.
33
34 You could potentially supply a file-based dict-like object, too.
35 This implementation deliberately avoid mapping.__iter__(),
36 which could be slow on a file-based mapping.
37
38 :param int capacity:
39 How many items this mapping will hold.
40 When you attempt to add new item into a full mapping,
41 it will automatically delete the item that is expiring soonest.
42
43 The default value is None, which means there is no capacity limit.
44
45 :param int expires_in:
46 How many seconds an item would expire and be purged from this mapping.
47 Also known as time-to-live (TTL).
48 You can also use :func:`~set()` to provide per-item expires_in value.
49
50 :param Lock lock:
51 A locking mechanism with context manager interface.
52 If no lock is provided, a threading.Lock will be used.
53 But you may want to supply a different lock,
54 if your customized mapping is being shared differently.
55 """
56 super(_ExpiringMapping, self).__init__(*args, **kwargs)
57 self._mapping = mapping if mapping is not None else {}
58 self._capacity = capacity
59 self._expires_in = expires_in
60 self._lock = Lock() if lock is None else lock
61
62 def _peek(self):
63 # Returns (sequence, timestamps) without triggering maintenance
64 return self._mapping.get(self._INDEX, ([], {}))
65
66 def _validate_key(self, key):
67 if key == self._INDEX:
68 raise ValueError("key {} is a reserved keyword in {}".format(
69 key, self.__class__.__name__))
70
71 def set(self, key, value, expires_in):
72 # This method's name was chosen so that it matches its cousin __setitem__(),
73 # and it also complements the counterpart get().
74 # The downside is such a name shadows the built-in type set in this file,
75 # but you can overcome that by defining a global alias for set.
76 """It sets the key-value pair into this mapping, with its per-item expires_in.
77
78 It will take O(logN) time, because it will run some maintenance.
79 This worse-than-constant time is acceptable, because in a cache scenario,
80 __setitem__() would only be called during a cache miss,
81 which would already incur an expensive target function call anyway.
82
83 By the way, most other methods of this mapping still have O(1) constant time.
84 """
85 with self._lock:
86 self._set(key, value, expires_in)
87
88 def _set(self, key, value, expires_in):
89 # This internal implementation powers both set() and __setitem__(),
90 # so that they don't depend on each other.
91 self._validate_key(key)
92 sequence, timestamps = self._peek()
93 self._maintenance(sequence, timestamps) # O(logN)
94 now = int(time.time())
95 expires_at = now + expires_in
96 entry = [expires_at, now, key]
97 is_new_item = key not in timestamps
98 is_beyond_capacity = self._capacity and len(timestamps) >= self._capacity
99 if is_new_item and is_beyond_capacity:
100 self._drop_indexed_entry(timestamps, heapq.heappushpop(sequence, entry))
101 else: # Simply add new entry. The old one would become a harmless orphan.
102 heapq.heappush(sequence, entry)
103 timestamps[key] = [expires_at, now] # It overwrites existing key, if any
104 self._mapping[key] = value
105 self._mapping[self._INDEX] = sequence, timestamps
106
107 def _maintenance(self, sequence, timestamps): # O(logN)
108 """It will modify input sequence and timestamps in-place"""
109 now = int(time.time())
110 while sequence: # Clean up expired items
111 expires_at, created_at, key = sequence[0]
112 if created_at <= now < expires_at: # Then all remaining items are fresh
113 break
114 self._drop_indexed_entry(timestamps, sequence[0]) # It could error out
115 heapq.heappop(sequence) # Only pop it after a successful _drop_indexed_entry()
116 while self._capacity is not None and len(timestamps) > self._capacity:
117 self._drop_indexed_entry(timestamps, sequence[0]) # It could error out
118 heapq.heappop(sequence) # Only pop it after a successful _drop_indexed_entry()
119
120 def _drop_indexed_entry(self, timestamps, entry):
121 """For an entry came from index, drop it from timestamps and self._mapping"""
122 expires_at, created_at, key = entry
123 if [expires_at, created_at] == timestamps.get(key): # So it is not an orphan
124 self._mapping.pop(key, None) # It could raise exception
125 timestamps.pop(key, None) # This would probably always succeed
126
127 def __setitem__(self, key, value):
128 """Implements the __setitem__().
129
130 Same characteristic as :func:`~set()`,
131 but use class-wide expires_in which was specified by :func:`~__init__()`.
132 """
133 if self._expires_in is None:
134 raise ValueError("Need a numeric value for expires_in during __init__()")
135 with self._lock:
136 self._set(key, value, self._expires_in)
137
138 def __getitem__(self, key): # O(1)
139 """If the item you requested already expires, KeyError will be raised."""
140 self._validate_key(key)
141 with self._lock:
142 # Skip self._maintenance(), because it would need O(logN) time
143 sequence, timestamps = self._peek()
144 expires_at, created_at = timestamps[key] # Would raise KeyError accordingly
145 now = int(time.time())
146 if not created_at <= now < expires_at:
147 self._mapping.pop(key, None)
148 timestamps.pop(key, None)
149 self._mapping[self._INDEX] = sequence, timestamps
150 raise KeyError("{} {}".format(
151 key,
152 "expired" if now >= expires_at else "created in the future?",
153 ))
154 return self._mapping[key] # O(1)
155
156 def __delitem__(self, key): # O(1)
157 """If the item you requested already expires, KeyError will be raised."""
158 self._validate_key(key)
159 with self._lock:
160 # Skip self._maintenance(), because it would need O(logN) time
161 self._mapping.pop(key, None) # O(1)
162 sequence, timestamps = self._peek()
163 del timestamps[key] # O(1)
164 self._mapping[self._INDEX] = sequence, timestamps
165
166 def __len__(self): # O(logN)
167 """Drop all expired items and return the remaining length"""
168 with self._lock:
169 sequence, timestamps = self._peek()
170 self._maintenance(sequence, timestamps) # O(logN)
171 self._mapping[self._INDEX] = sequence, timestamps
172 return len(timestamps) # Faster than iter(self._mapping) when it is on disk
173
174 def __iter__(self):
175 """Drop all expired items and return an iterator of the remaining items"""
176 with self._lock:
177 sequence, timestamps = self._peek()
178 self._maintenance(sequence, timestamps) # O(logN)
179 self._mapping[self._INDEX] = sequence, timestamps
180 return iter(timestamps) # Faster than iter(self._mapping) when it is on disk
181
182
183class _IndividualCache(object):
184 # The code structure below can decorate both function and method.
185 # It is inspired by https://stackoverflow.com/a/9417088
186 # We may potentially switch to build upon
187 # https://github.com/micheles/decorator/blob/master/docs/documentation.md#statement-of-the-problem
188 def __init__(self, mapping=None, key_maker=None, expires_in=None):
189 """Constructs a cache decorator that allows item-by-item control on
190 how to cache the return value of the decorated function.
191
192 :param MutableMapping mapping:
193 The cached items will be stored inside.
194 You'd want to use a ExpiringMapping
195 if you plan to utilize the ``expires_in`` behavior.
196
197 If nothing is provided, an in-memory dict will be used,
198 but it will provide no expiry functionality.
199
200 .. note::
201
202 When using this class as a decorator,
203 your mapping needs to be available at "compile" time,
204 so it would typically be a global-, module- or class-level mapping::
205
206 module_mapping = {}
207
208 @IndividualCache(mapping=module_mapping, ...)
209 def foo():
210 ...
211
212 If you want to use a mapping available only at run-time,
213 you have to manually decorate your function at run-time, too::
214
215 def foo():
216 ...
217
218 def bar(runtime_mapping):
219 foo = IndividualCache(mapping=runtime_mapping...)(foo)
220
221 :param callable key_maker:
222 A callable which should have signature as
223 ``lambda function, args, kwargs: "return a string as key"``.
224
225 If key_maker happens to return ``None``, the cache will be bypassed,
226 the underlying function will be invoked directly,
227 and the invoke result will not be cached either.
228
229 :param callable expires_in:
230 The default value is ``None``,
231 which means the content being cached has no per-item expiry,
232 and will subject to the underlying mapping's global expiry time.
233
234 It can be an integer indicating
235 how many seconds the result will be cached.
236 In particular, if the value is 0,
237 it means the result expires after zero second (i.e. immediately),
238 therefore the result will *not* be cached.
239 (Mind the difference between ``expires_in=0`` and ``expires_in=None``.)
240
241 Or it can be a callable with the signature as
242 ``lambda function=function, args=args, kwargs=kwargs, result=result: 123``
243 to calculate the expiry on the fly.
244 Its return value will be interpreted in the same way as above.
245 """
246 self._mapping = mapping if mapping is not None else {}
247 self._key_maker = key_maker or (lambda function, args, kwargs: (
248 function, # This default implementation uses function as part of key,
249 # so that the cache is partitioned by function.
250 # However, you could have many functions to use same namespace,
251 # so different decorators could share same cache.
252 args,
253 tuple(kwargs.items()), # raw kwargs is not hashable
254 ))
255 self._expires_in = expires_in
256
257 def __call__(self, function):
258
259 @wraps(function)
260 def wrapper(*args, **kwargs):
261 key = self._key_maker(function, args, kwargs)
262 if key is None: # Then bypass the cache
263 return function(*args, **kwargs)
264
265 now = int(time.time())
266 try:
267 return self._mapping[key]
268 except KeyError:
269 # We choose to NOT call function(...) in this block, otherwise
270 # potential exception from function(...) would become a confusing
271 # "During handling of the above exception, another exception occurred"
272 pass
273 value = function(*args, **kwargs)
274
275 expires_in = self._expires_in(
276 function=function,
277 args=args,
278 kwargs=kwargs,
279 result=value,
280 ) if callable(self._expires_in) else self._expires_in
281 if expires_in == 0:
282 return value
283 if expires_in is None:
284 self._mapping[key] = value
285 else:
286 self._mapping.set(key, value, expires_in)
287 return value
288
289 return wrapper
290