1# -*- coding: utf-8 -*-
2
3# -----------------------------------------------------------------------------
4# Copyright (c) 2015, Enthought, Inc.
5# All rights reserved.
6#
7# This software is provided without warranty under the terms of the BSD
8# license included in enthought/LICENSE.txt and may be redistributed only
9# under the conditions described in the aforementioned license.  The license
10# is also available online at http://www.enthought.com/licenses/BSD.txt
11# Thanks for using Enthought open source!
12#
13# Author: Enthought, Inc.
14# -----------------------------------------------------------------------------
15
16from threading import RLock
17
18try:
19    from collections import OrderedDict
20except ImportError:
21    from ordereddict import OrderedDict
22
23
24from traits.api import Callable, Event, HasStrictTraits, Instance, Int
25
26
27class LRUCache(HasStrictTraits):
28    """ A least-recently used cache.
29
30    Items older than `size()` accesses are dropped from the cache.
31
32    """
33
34    size = Int
35
36    # Called with the key and value that was dropped from the cache
37    cache_drop_callback = Callable
38
39    # This event contains the set of cached cell keys whenever it changes
40    updated = Event()
41
42    _lock = Instance(RLock, args=())
43
44    _cache = Instance(OrderedDict)
45
46    def __init__(self, size, **traits):
47        self.size = size
48        self._initialize_cache()
49        super(LRUCache, self).__init__(**traits)
50
51    def _initialize_cache(self):
52        with self._lock:
53            if self._cache is None:
54                self._cache = OrderedDict()
55            else:
56                self._cache.clear()
57
58    def _renew(self, key):
59        with self._lock:
60            r = self._cache.pop(key)
61            self._cache[key] = r
62        return r
63
64    # -------------------------------------------------------------------------
65    # LRUCache interface
66    # -------------------------------------------------------------------------
67
68    def __contains__(self, key):
69        with self._lock:
70            return key in self._cache
71
72    def __len__(self):
73        with self._lock:
74            return len(self._cache)
75
76    def __getitem__(self, key):
77        with self._lock:
78            return self._renew(key)
79
80    def __setitem__(self, key, result):
81        try:
82            dropped = None
83            with self._lock:
84                self._cache[key] = result
85                self._renew(key)
86                if self.size < len(self._cache):
87                    dropped = self._cache.popitem(last=False)
88            if dropped and self.cache_drop_callback is not None:
89                self.cache_drop_callback(*dropped)
90        finally:
91            self.updated = list(self.keys())
92
93    def get(self, key, default=None):
94        try:
95            return self[key]
96        except KeyError:
97            return default
98
99    def items(self):
100        with self._lock:
101            return list(self._cache.items())
102
103    def keys(self):
104        with self._lock:
105            return list(self._cache.keys())
106
107    def values(self):
108        with self._lock:
109            return list(self._cache.values())
110
111    def clear(self):
112        with self._lock:
113            self._initialize_cache()
114        self.updated = []
115