You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

78 lines
2.6 KiB

  1. # Memoize a function's return value with a least-recently-used cache
  2. # Based on:
  3. # http://code.activestate.com/recipes/498245-lru-and-lfu-cache-decorators/
  4. # with added 'destructor' functionality.
  5. import collections
  6. import decorator
  7. import warnings
  8. def lru_cache(size = 10, onremove = None, keys = slice(None)):
  9. """Least-recently-used cache decorator.
  10. @lru_cache(size = 10, onevict = None)
  11. def f(...):
  12. pass
  13. Given a function and arguments, memoize its return value. Up to
  14. 'size' elements are cached. 'keys' is a slice object that
  15. represents which arguments are used as the cache key.
  16. When evicting a value from the cache, call the function
  17. 'onremove' with the value that's being evicted.
  18. Call f.cache_remove(...) to evict the cache entry with the given
  19. arguments. Call f.cache_remove_all() to evict all entries.
  20. f.cache_hits and f.cache_misses give statistics.
  21. """
  22. def decorate(func):
  23. cache = collections.OrderedDict() # order: least- to most-recent
  24. def evict(value):
  25. if onremove:
  26. onremove(value)
  27. def wrapper(orig, *args, **kwargs):
  28. if kwargs:
  29. raise NotImplementedError("kwargs not supported")
  30. key = args[keys]
  31. try:
  32. value = cache.pop(key)
  33. orig.cache_hits += 1
  34. except KeyError:
  35. value = orig(*args)
  36. orig.cache_misses += 1
  37. if len(cache) >= size:
  38. evict(cache.popitem(0)[1]) # evict LRU cache entry
  39. cache[key] = value # (re-)insert this key at end
  40. return value
  41. def cache_remove(*args):
  42. """Remove the described key from this cache, if present."""
  43. key = args
  44. if key in cache:
  45. evict(cache.pop(key))
  46. else:
  47. if len(cache) > 0 and len(args) != len(cache.iterkeys().next()):
  48. raise KeyError("trying to remove from LRU cache, but "
  49. "number of arguments doesn't match the "
  50. "cache key length")
  51. def cache_remove_all():
  52. for key in cache:
  53. evict(cache.pop(key))
  54. def cache_info():
  55. return (func.cache_hits, func.cache_misses)
  56. new = decorator.decorator(wrapper, func)
  57. func.cache_hits = 0
  58. func.cache_misses = 0
  59. new.cache_info = cache_info
  60. new.cache_remove = cache_remove
  61. new.cache_remove_all = cache_remove_all
  62. return new
  63. return decorate