1"""Extensible memoizing collections and decorators."""
2
3__all__ = (
4    "Cache",
5    "FIFOCache",
6    "LFUCache",
7    "LRUCache",
8    "MRUCache",
9    "RRCache",
10    "TTLCache",
11    "cached",
12    "cachedmethod",
13)
14
15__version__ = "4.2.4"
16
17import collections
18import collections.abc
19import functools
20import random
21import time
22
23from .keys import hashkey
24
25
26class _DefaultSize:
27
28    __slots__ = ()
29
30    def __getitem__(self, _):
31        return 1
32
33    def __setitem__(self, _, value):
34        assert value == 1
35
36    def pop(self, _):
37        return 1
38
39
40class Cache(collections.abc.MutableMapping):
41    """Mutable mapping to serve as a simple cache or cache base class."""
42
43    __marker = object()
44
45    __size = _DefaultSize()
46
47    def __init__(self, maxsize, getsizeof=None):
48        if getsizeof:
49            self.getsizeof = getsizeof
50        if self.getsizeof is not Cache.getsizeof:
51            self.__size = dict()
52        self.__data = dict()
53        self.__currsize = 0
54        self.__maxsize = maxsize
55
56    def __repr__(self):
57        return "%s(%r, maxsize=%r, currsize=%r)" % (
58            self.__class__.__name__,
59            list(self.__data.items()),
60            self.__maxsize,
61            self.__currsize,
62        )
63
64    def __getitem__(self, key):
65        try:
66            return self.__data[key]
67        except KeyError:
68            return self.__missing__(key)
69
70    def __setitem__(self, key, value):
71        maxsize = self.__maxsize
72        size = self.getsizeof(value)
73        if size > maxsize:
74            raise ValueError("value too large")
75        if key not in self.__data or self.__size[key] < size:
76            while self.__currsize + size > maxsize:
77                self.popitem()
78        if key in self.__data:
79            diffsize = size - self.__size[key]
80        else:
81            diffsize = size
82        self.__data[key] = value
83        self.__size[key] = size
84        self.__currsize += diffsize
85
86    def __delitem__(self, key):
87        size = self.__size.pop(key)
88        del self.__data[key]
89        self.__currsize -= size
90
91    def __contains__(self, key):
92        return key in self.__data
93
94    def __missing__(self, key):
95        raise KeyError(key)
96
97    def __iter__(self):
98        return iter(self.__data)
99
100    def __len__(self):
101        return len(self.__data)
102
103    def get(self, key, default=None):
104        if key in self:
105            return self[key]
106        else:
107            return default
108
109    def pop(self, key, default=__marker):
110        if key in self:
111            value = self[key]
112            del self[key]
113        elif default is self.__marker:
114            raise KeyError(key)
115        else:
116            value = default
117        return value
118
119    def setdefault(self, key, default=None):
120        if key in self:
121            value = self[key]
122        else:
123            self[key] = value = default
124        return value
125
126    @property
127    def maxsize(self):
128        """The maximum size of the cache."""
129        return self.__maxsize
130
131    @property
132    def currsize(self):
133        """The current size of the cache."""
134        return self.__currsize
135
136    @staticmethod
137    def getsizeof(value):
138        """Return the size of a cache element's value."""
139        return 1
140
141
142class FIFOCache(Cache):
143    """First In First Out (FIFO) cache implementation."""
144
145    def __init__(self, maxsize, getsizeof=None):
146        Cache.__init__(self, maxsize, getsizeof)
147        self.__order = collections.OrderedDict()
148
149    def __setitem__(self, key, value, cache_setitem=Cache.__setitem__):
150        cache_setitem(self, key, value)
151        try:
152            self.__order.move_to_end(key)
153        except KeyError:
154            self.__order[key] = None
155
156    def __delitem__(self, key, cache_delitem=Cache.__delitem__):
157        cache_delitem(self, key)
158        del self.__order[key]
159
160    def popitem(self):
161        """Remove and return the `(key, value)` pair first inserted."""
162        try:
163            key = next(iter(self.__order))
164        except StopIteration:
165            raise KeyError("%s is empty" % type(self).__name__) from None
166        else:
167            return (key, self.pop(key))
168
169
170class LFUCache(Cache):
171    """Least Frequently Used (LFU) cache implementation."""
172
173    def __init__(self, maxsize, getsizeof=None):
174        Cache.__init__(self, maxsize, getsizeof)
175        self.__counter = collections.Counter()
176
177    def __getitem__(self, key, cache_getitem=Cache.__getitem__):
178        value = cache_getitem(self, key)
179        if key in self:  # __missing__ may not store item
180            self.__counter[key] -= 1
181        return value
182
183    def __setitem__(self, key, value, cache_setitem=Cache.__setitem__):
184        cache_setitem(self, key, value)
185        self.__counter[key] -= 1
186
187    def __delitem__(self, key, cache_delitem=Cache.__delitem__):
188        cache_delitem(self, key)
189        del self.__counter[key]
190
191    def popitem(self):
192        """Remove and return the `(key, value)` pair least frequently used."""
193        try:
194            ((key, _),) = self.__counter.most_common(1)
195        except ValueError:
196            raise KeyError("%s is empty" % type(self).__name__) from None
197        else:
198            return (key, self.pop(key))
199
200
201class LRUCache(Cache):
202    """Least Recently Used (LRU) cache implementation."""
203
204    def __init__(self, maxsize, getsizeof=None):
205        Cache.__init__(self, maxsize, getsizeof)
206        self.__order = collections.OrderedDict()
207
208    def __getitem__(self, key, cache_getitem=Cache.__getitem__):
209        value = cache_getitem(self, key)
210        if key in self:  # __missing__ may not store item
211            self.__update(key)
212        return value
213
214    def __setitem__(self, key, value, cache_setitem=Cache.__setitem__):
215        cache_setitem(self, key, value)
216        self.__update(key)
217
218    def __delitem__(self, key, cache_delitem=Cache.__delitem__):
219        cache_delitem(self, key)
220        del self.__order[key]
221
222    def popitem(self):
223        """Remove and return the `(key, value)` pair least recently used."""
224        try:
225            key = next(iter(self.__order))
226        except StopIteration:
227            raise KeyError("%s is empty" % type(self).__name__) from None
228        else:
229            return (key, self.pop(key))
230
231    def __update(self, key):
232        try:
233            self.__order.move_to_end(key)
234        except KeyError:
235            self.__order[key] = None
236
237
238class MRUCache(Cache):
239    """Most Recently Used (MRU) cache implementation."""
240
241    def __init__(self, maxsize, getsizeof=None):
242        Cache.__init__(self, maxsize, getsizeof)
243        self.__order = collections.OrderedDict()
244
245    def __getitem__(self, key, cache_getitem=Cache.__getitem__):
246        value = cache_getitem(self, key)
247        if key in self:  # __missing__ may not store item
248            self.__update(key)
249        return value
250
251    def __setitem__(self, key, value, cache_setitem=Cache.__setitem__):
252        cache_setitem(self, key, value)
253        self.__update(key)
254
255    def __delitem__(self, key, cache_delitem=Cache.__delitem__):
256        cache_delitem(self, key)
257        del self.__order[key]
258
259    def popitem(self):
260        """Remove and return the `(key, value)` pair most recently used."""
261        try:
262            key = next(iter(self.__order))
263        except StopIteration:
264            raise KeyError("%s is empty" % type(self).__name__) from None
265        else:
266            return (key, self.pop(key))
267
268    def __update(self, key):
269        try:
270            self.__order.move_to_end(key, last=False)
271        except KeyError:
272            self.__order[key] = None
273
274
275class RRCache(Cache):
276    """Random Replacement (RR) cache implementation."""
277
278    def __init__(self, maxsize, choice=random.choice, getsizeof=None):
279        Cache.__init__(self, maxsize, getsizeof)
280        self.__choice = choice
281
282    @property
283    def choice(self):
284        """The `choice` function used by the cache."""
285        return self.__choice
286
287    def popitem(self):
288        """Remove and return a random `(key, value)` pair."""
289        try:
290            key = self.__choice(list(self))
291        except IndexError:
292            raise KeyError("%s is empty" % type(self).__name__) from None
293        else:
294            return (key, self.pop(key))
295
296
297class _Timer:
298    def __init__(self, timer):
299        self.__timer = timer
300        self.__nesting = 0
301
302    def __call__(self):
303        if self.__nesting == 0:
304            return self.__timer()
305        else:
306            return self.__time
307
308    def __enter__(self):
309        if self.__nesting == 0:
310            self.__time = time = self.__timer()
311        else:
312            time = self.__time
313        self.__nesting += 1
314        return time
315
316    def __exit__(self, *exc):
317        self.__nesting -= 1
318
319    def __reduce__(self):
320        return _Timer, (self.__timer,)
321
322    def __getattr__(self, name):
323        return getattr(self.__timer, name)
324
325
326class _Link:
327
328    __slots__ = ("key", "expire", "next", "prev")
329
330    def __init__(self, key=None, expire=None):
331        self.key = key
332        self.expire = expire
333
334    def __reduce__(self):
335        return _Link, (self.key, self.expire)
336
337    def unlink(self):
338        next = self.next
339        prev = self.prev
340        prev.next = next
341        next.prev = prev
342
343
344class TTLCache(Cache):
345    """LRU Cache implementation with per-item time-to-live (TTL) value."""
346
347    def __init__(self, maxsize, ttl, timer=time.monotonic, getsizeof=None):
348        Cache.__init__(self, maxsize, getsizeof)
349        self.__root = root = _Link()
350        root.prev = root.next = root
351        self.__links = collections.OrderedDict()
352        self.__timer = _Timer(timer)
353        self.__ttl = ttl
354
355    def __contains__(self, key):
356        try:
357            link = self.__links[key]  # no reordering
358        except KeyError:
359            return False
360        else:
361            return not (link.expire < self.__timer())
362
363    def __getitem__(self, key, cache_getitem=Cache.__getitem__):
364        try:
365            link = self.__getlink(key)
366        except KeyError:
367            expired = False
368        else:
369            expired = link.expire < self.__timer()
370        if expired:
371            return self.__missing__(key)
372        else:
373            return cache_getitem(self, key)
374
375    def __setitem__(self, key, value, cache_setitem=Cache.__setitem__):
376        with self.__timer as time:
377            self.expire(time)
378            cache_setitem(self, key, value)
379        try:
380            link = self.__getlink(key)
381        except KeyError:
382            self.__links[key] = link = _Link(key)
383        else:
384            link.unlink()
385        link.expire = time + self.__ttl
386        link.next = root = self.__root
387        link.prev = prev = root.prev
388        prev.next = root.prev = link
389
390    def __delitem__(self, key, cache_delitem=Cache.__delitem__):
391        cache_delitem(self, key)
392        link = self.__links.pop(key)
393        link.unlink()
394        if link.expire < self.__timer():
395            raise KeyError(key)
396
397    def __iter__(self):
398        root = self.__root
399        curr = root.next
400        while curr is not root:
401            # "freeze" time for iterator access
402            with self.__timer as time:
403                if not (curr.expire < time):
404                    yield curr.key
405            curr = curr.next
406
407    def __len__(self):
408        root = self.__root
409        curr = root.next
410        time = self.__timer()
411        count = len(self.__links)
412        while curr is not root and curr.expire < time:
413            count -= 1
414            curr = curr.next
415        return count
416
417    def __setstate__(self, state):
418        self.__dict__.update(state)
419        root = self.__root
420        root.prev = root.next = root
421        for link in sorted(self.__links.values(), key=lambda obj: obj.expire):
422            link.next = root
423            link.prev = prev = root.prev
424            prev.next = root.prev = link
425        self.expire(self.__timer())
426
427    def __repr__(self, cache_repr=Cache.__repr__):
428        with self.__timer as time:
429            self.expire(time)
430            return cache_repr(self)
431
432    @property
433    def currsize(self):
434        with self.__timer as time:
435            self.expire(time)
436            return super().currsize
437
438    @property
439    def timer(self):
440        """The timer function used by the cache."""
441        return self.__timer
442
443    @property
444    def ttl(self):
445        """The time-to-live value of the cache's items."""
446        return self.__ttl
447
448    def expire(self, time=None):
449        """Remove expired items from the cache."""
450        if time is None:
451            time = self.__timer()
452        root = self.__root
453        curr = root.next
454        links = self.__links
455        cache_delitem = Cache.__delitem__
456        while curr is not root and curr.expire < time:
457            cache_delitem(self, curr.key)
458            del links[curr.key]
459            next = curr.next
460            curr.unlink()
461            curr = next
462
463    def clear(self):
464        with self.__timer as time:
465            self.expire(time)
466            Cache.clear(self)
467
468    def get(self, *args, **kwargs):
469        with self.__timer:
470            return Cache.get(self, *args, **kwargs)
471
472    def pop(self, *args, **kwargs):
473        with self.__timer:
474            return Cache.pop(self, *args, **kwargs)
475
476    def setdefault(self, *args, **kwargs):
477        with self.__timer:
478            return Cache.setdefault(self, *args, **kwargs)
479
480    def popitem(self):
481        """Remove and return the `(key, value)` pair least recently used that
482        has not already expired.
483
484        """
485        with self.__timer as time:
486            self.expire(time)
487            try:
488                key = next(iter(self.__links))
489            except StopIteration:
490                raise KeyError("%s is empty" % type(self).__name__) from None
491            else:
492                return (key, self.pop(key))
493
494    def __getlink(self, key):
495        value = self.__links[key]
496        self.__links.move_to_end(key)
497        return value
498
499
500def cached(cache, key=hashkey, lock=None):
501    """Decorator to wrap a function with a memoizing callable that saves
502    results in a cache.
503
504    """
505
506    def decorator(func):
507        if cache is None:
508
509            def wrapper(*args, **kwargs):
510                return func(*args, **kwargs)
511
512        elif lock is None:
513
514            def wrapper(*args, **kwargs):
515                k = key(*args, **kwargs)
516                try:
517                    return cache[k]
518                except KeyError:
519                    pass  # key not found
520                v = func(*args, **kwargs)
521                try:
522                    cache[k] = v
523                except ValueError:
524                    pass  # value too large
525                return v
526
527        else:
528
529            def wrapper(*args, **kwargs):
530                k = key(*args, **kwargs)
531                try:
532                    with lock:
533                        return cache[k]
534                except KeyError:
535                    pass  # key not found
536                v = func(*args, **kwargs)
537                # in case of a race, prefer the item already in the cache
538                try:
539                    with lock:
540                        return cache.setdefault(k, v)
541                except ValueError:
542                    return v  # value too large
543
544        return functools.update_wrapper(wrapper, func)
545
546    return decorator
547
548
549def cachedmethod(cache, key=hashkey, lock=None):
550    """Decorator to wrap a class or instance method with a memoizing
551    callable that saves results in a cache.
552
553    """
554
555    def decorator(method):
556        if lock is None:
557
558            def wrapper(self, *args, **kwargs):
559                c = cache(self)
560                if c is None:
561                    return method(self, *args, **kwargs)
562                k = key(*args, **kwargs)
563                try:
564                    return c[k]
565                except KeyError:
566                    pass  # key not found
567                v = method(self, *args, **kwargs)
568                try:
569                    c[k] = v
570                except ValueError:
571                    pass  # value too large
572                return v
573
574        else:
575
576            def wrapper(self, *args, **kwargs):
577                c = cache(self)
578                if c is None:
579                    return method(self, *args, **kwargs)
580                k = key(*args, **kwargs)
581                try:
582                    with lock(self):
583                        return c[k]
584                except KeyError:
585                    pass  # key not found
586                v = method(self, *args, **kwargs)
587                # in case of a race, prefer the item already in the cache
588                try:
589                    with lock(self):
590                        return c.setdefault(k, v)
591                except ValueError:
592                    return v  # value too large
593
594        return functools.update_wrapper(wrapper, method)
595
596    return decorator
597