- # Memoize a function's return value with a least-recently-used cache
- # Based on:
- # http://code.activestate.com/recipes/498245-lru-and-lfu-cache-decorators/
- # with added 'destructor' functionality.
-
- import collections
- import decorator
- import warnings
-
- def lru_cache(size = 10, onremove = None, keys = slice(None)):
- """Least-recently-used cache decorator.
-
- @lru_cache(size = 10, onevict = None)
- def f(...):
- pass
-
- Given a function and arguments, memoize its return value. Up to
- 'size' elements are cached. 'keys' is a slice object that
- represents which arguments are used as the cache key.
-
- When evicting a value from the cache, call the function
- 'onremove' with the value that's being evicted.
-
- Call f.cache_remove(...) to evict the cache entry with the given
- arguments. Call f.cache_remove_all() to evict all entries.
- f.cache_hits and f.cache_misses give statistics.
- """
-
- def decorate(func):
- cache = collections.OrderedDict() # order: least- to most-recent
-
- def evict(value):
- if onremove:
- onremove(value)
-
- def wrapper(orig, *args, **kwargs):
- if kwargs:
- raise NotImplementedError("kwargs not supported")
- key = args[keys]
- try:
- value = cache.pop(key)
- orig.cache_hits += 1
- except KeyError:
- value = orig(*args)
- orig.cache_misses += 1
- if len(cache) >= size:
- evict(cache.popitem(0)[1]) # evict LRU cache entry
- cache[key] = value # (re-)insert this key at end
- return value
-
- def cache_remove(*args):
- """Remove the described key from this cache, if present."""
- key = args
- if key in cache:
- evict(cache.pop(key))
- else:
- if len(cache) > 0 and len(args) != len(cache.iterkeys().next()):
- raise KeyError("trying to remove from LRU cache, but "
- "number of arguments doesn't match the "
- "cache key length")
-
- def cache_remove_all():
- for key in cache:
- evict(cache.pop(key))
-
- def cache_info():
- return (func.cache_hits, func.cache_misses)
-
- new = decorator.decorator(wrapper, func)
- func.cache_hits = 0
- func.cache_misses = 0
- new.cache_info = cache_info
- new.cache_remove = cache_remove
- new.cache_remove_all = cache_remove_all
- return new
-
- return decorate
|