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.

147 lines
3.9 KiB

  1. """`functools.lru_cache` compatible memoizing function decorators."""
  2. import collections
  3. import functools
  4. import math
  5. import random
  6. import time
  7. try:
  8. from threading import RLock
  9. except ImportError: # pragma: no cover
  10. from dummy_threading import RLock
  11. from . import keys
  12. from .lfu import LFUCache
  13. from .lru import LRUCache
  14. from .rr import RRCache
  15. from .ttl import TTLCache
  16. __all__ = ('lfu_cache', 'lru_cache', 'rr_cache', 'ttl_cache')
  17. _CacheInfo = collections.namedtuple('CacheInfo', [
  18. 'hits', 'misses', 'maxsize', 'currsize'
  19. ])
  20. class _UnboundCache(dict):
  21. @property
  22. def maxsize(self):
  23. return None
  24. @property
  25. def currsize(self):
  26. return len(self)
  27. class _UnboundTTLCache(TTLCache):
  28. def __init__(self, ttl, timer):
  29. TTLCache.__init__(self, math.inf, ttl, timer)
  30. @property
  31. def maxsize(self):
  32. return None
  33. def _cache(cache, typed):
  34. maxsize = cache.maxsize
  35. def decorator(func):
  36. key = keys.typedkey if typed else keys.hashkey
  37. lock = RLock()
  38. stats = [0, 0]
  39. def wrapper(*args, **kwargs):
  40. k = key(*args, **kwargs)
  41. with lock:
  42. try:
  43. v = cache[k]
  44. stats[0] += 1
  45. return v
  46. except KeyError:
  47. stats[1] += 1
  48. v = func(*args, **kwargs)
  49. try:
  50. with lock:
  51. cache[k] = v
  52. except ValueError:
  53. pass # value too large
  54. return v
  55. def cache_info():
  56. with lock:
  57. hits, misses = stats
  58. maxsize = cache.maxsize
  59. currsize = cache.currsize
  60. return _CacheInfo(hits, misses, maxsize, currsize)
  61. def cache_clear():
  62. with lock:
  63. try:
  64. cache.clear()
  65. finally:
  66. stats[:] = [0, 0]
  67. wrapper.cache_info = cache_info
  68. wrapper.cache_clear = cache_clear
  69. wrapper.cache_parameters = lambda: {'maxsize': maxsize, 'typed': typed}
  70. functools.update_wrapper(wrapper, func)
  71. return wrapper
  72. return decorator
  73. def lfu_cache(maxsize=128, typed=False):
  74. """Decorator to wrap a function with a memoizing callable that saves
  75. up to `maxsize` results based on a Least Frequently Used (LFU)
  76. algorithm.
  77. """
  78. if maxsize is None:
  79. return _cache(_UnboundCache(), typed)
  80. elif callable(maxsize):
  81. return _cache(LFUCache(128), typed)(maxsize)
  82. else:
  83. return _cache(LFUCache(maxsize), typed)
  84. def lru_cache(maxsize=128, typed=False):
  85. """Decorator to wrap a function with a memoizing callable that saves
  86. up to `maxsize` results based on a Least Recently Used (LRU)
  87. algorithm.
  88. """
  89. if maxsize is None:
  90. return _cache(_UnboundCache(), typed)
  91. elif callable(maxsize):
  92. return _cache(LRUCache(128), typed)(maxsize)
  93. else:
  94. return _cache(LRUCache(maxsize), typed)
  95. def rr_cache(maxsize=128, choice=random.choice, typed=False):
  96. """Decorator to wrap a function with a memoizing callable that saves
  97. up to `maxsize` results based on a Random Replacement (RR)
  98. algorithm.
  99. """
  100. if maxsize is None:
  101. return _cache(_UnboundCache(), typed)
  102. elif callable(maxsize):
  103. return _cache(RRCache(128, choice), typed)(maxsize)
  104. else:
  105. return _cache(RRCache(maxsize, choice), typed)
  106. def ttl_cache(maxsize=128, ttl=600, timer=time.monotonic, typed=False):
  107. """Decorator to wrap a function with a memoizing callable that saves
  108. up to `maxsize` results based on a Least Recently Used (LRU)
  109. algorithm with a per-item time-to-live (TTL) value.
  110. """
  111. if maxsize is None:
  112. return _cache(_UnboundTTLCache(ttl, timer), typed)
  113. elif callable(maxsize):
  114. return _cache(TTLCache(128, ttl, timer), typed)(maxsize)
  115. else:
  116. return _cache(TTLCache(maxsize, ttl, timer), typed)