bazarr/libs/wraptor/decorators/memoize.py

71 lines
2.6 KiB
Python

from functools import wraps
import time
from hashlib import md5
import threading
class memoize(object):
""" Memoize the results of a function. Supports an optional timeout
for automatic cache expiration.
If the optional manual_flush argument is True, a function called
"flush_cache" will be added to the wrapped function. When
called, it will remove all the timed out values from the cache.
If you use this decorator as a class method, you must specify
instance_method=True otherwise you will have a single shared
cache for every instance of your class.
This decorator is thread safe.
"""
def __init__(self, timeout=None, manual_flush=False, instance_method=False):
self.timeout = timeout
self.manual_flush = manual_flush
self.instance_method = instance_method
self.cache = {}
self.cache_lock = threading.RLock()
def __call__(self, fn):
if self.instance_method:
@wraps(fn)
def rewrite_instance_method(instance, *args, **kwargs):
# the first time we are called we overwrite the method
# on the class instance with a new memoize instance
if hasattr(instance, fn.__name__):
bound_fn = fn.__get__(instance, instance.__class__)
new_memoizer = memoize(self.timeout, self.manual_flush)(bound_fn)
setattr(instance, fn.__name__, new_memoizer)
return getattr(instance, fn.__name__)(*args, **kwargs)
return rewrite_instance_method
def flush_cache():
with self.cache_lock:
for key in self.cache.keys():
if (time.time() - self.cache[key][1]) > self.timeout:
del(self.cache[key])
@wraps(fn)
def wrapped(*args, **kwargs):
kw = kwargs.items()
kw.sort()
key_str = repr((args, kw))
key = md5(key_str).hexdigest()
with self.cache_lock:
try:
result, cache_time = self.cache[key]
if self.timeout is not None and (time.time() - cache_time) > self.timeout:
raise KeyError
except KeyError:
result, _ = self.cache[key] = (fn(*args, **kwargs), time.time())
if not self.manual_flush and self.timeout is not None:
flush_cache()
return result
if self.manual_flush:
wrapped.flush_cache = flush_cache
return wrapped