Simpler and faster LRUCache implementation

This commit is contained in:
Jonas Borgström 2012-12-06 22:57:55 +01:00
parent f28933254a
commit e80e600d41
1 changed files with 25 additions and 66 deletions

View File

@ -3,78 +3,37 @@ from heapq import heappush, heapify, heapreplace, heappop
import unittest
class LRUCache(DictMixin):
"""Heap queue based Least Recently Used Cache implementation
"""
class LRUCache(dict):
class Node(object):
"""Internal cache node
"""
__slots__ = ('key', 'value', 't')
def __init__(self, key, value, t):
self.key = key
self.value = value
self.t = t
def __cmp__(self, other):
return cmp(self.t, other.t)
def __init__(self, size):
self.size = size
self._t = 0
self.clear()
def clear(self):
self._heap = []
self._dict = {}
def __init__(self, capacity):
super(LRUCache, self).__init__()
self._lru = []
self._capacity = capacity
def __setitem__(self, key, value):
self._t += 1
try:
node = self._dict[key]
node.value = value
node.t = self._t
heapify(self._heap)
except KeyError:
node = self.Node(key, value, self._t)
self._dict[key] = node
if len(self) < self.size:
heappush(self._heap, node)
else:
old = heapreplace(self._heap, node)
del self._dict[old.key]
self._lru.remove(key)
except ValueError:
pass
self._lru.append(key)
while len(self._lru) > self._capacity:
del self[self._lru[0]]
return super(LRUCache, self).__setitem__(key, value)
def __getitem__(self, key):
node = self._dict[key]
self[key] = node.value
return node.value
try:
self._lru.remove(key)
self._lru.append(key)
except ValueError:
pass
return super(LRUCache, self).__getitem__(key)
def __delitem__(self, key):
node = self._dict[key]
del self._dict[key]
self._heap.remove(node)
heapify(self._heap)
def __iter__(self):
copy = self._heap[:]
while copy:
yield heappop(copy).key
def iteritems(self):
copy = self._heap[:]
while copy:
node = heappop(copy)
yield node.key, node.value
def keys(self):
return self._dict.keys()
def __contains__(self, key):
return key in self._dict
def __len__(self):
return len(self._heap)
try:
self._lru.remove(key)
except ValueError:
pass
return super(LRUCache, self).__delitem__(key)
class LRUCacheTestCase(unittest.TestCase):
@ -85,8 +44,8 @@ class LRUCacheTestCase(unittest.TestCase):
for i, x in enumerate('abc'):
c[x] = i
self.assertEqual(len(c), 2)
self.assertEqual(list(c), ['b', 'c'])
self.assertEqual(list(c.iteritems()), [('b', 1), ('c', 2)])
self.assertEqual(set(c), set(['b', 'c']))
self.assertEqual(set(c.iteritems()), set([('b', 1), ('c', 2)]))
self.assertEqual(False, 'a' in c)
self.assertEqual(True, 'b' in c)
self.assertRaises(KeyError, lambda: c['a'])