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.
 
 
 

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