Support args & kwargs for memoize
This is needed to cache results based on arg/kwarg values.
Haven't studied the perf impact of this yet
#!/usr/bin/env python3
import functools
+import pickle
-def add_to_cache(obj, name, val):
+def add_to_cache(obj, name, val, *args, **kwargs):
+ """Add a result to the cache of an object."""
+ return _add_to_cache(obj, name, *args, pickle.dumps(kwargs))
+
+
+def _add_to_cache(obj, name, val, *args, kwargs_pkl):
"""Add a result to the cache of an object."""
if not hasattr(obj, "_memoize_cache"):
- obj._memoize_cache = dict()
- obj._memoize_cache[name] = val
+ obj._memoize_cache = {}
+ obj._memoize_cache[(name, args, kwargs_pkl)] = val
return obj
-def get_from_cache(obj, name):
+def _get_from_cache(obj, name, *args, kwargs_pkl):
"""Get an item from the cache."""
- if not is_in_cache(obj, name):
+ if not _is_in_cache(obj, name, *args, kwargs_pkl=kwargs_pkl):
raise RuntimeError("Object does not have item {} stored in cache.".format(name))
- return obj._memoize_cache[name]
+ return obj._memoize_cache[(name, args, kwargs_pkl)]
-def is_in_cache(obj, name):
- return hasattr(obj, "_memoize_cache") and name in obj._memoize_cache
+def _is_in_cache(obj, name, *args, kwargs_pkl):
+ return hasattr(obj, "_memoize_cache") and (name, args, kwargs_pkl) in obj._memoize_cache
def cached(method=None, name=None):
@@ -30,15 +36,9 @@ def cached(method=None, name=None):
@functools.wraps(method)
def g(self, *args, **kwargs):
cache_name = name if name is not None else method
- if not is_in_cache(self, cache_name):
- add_to_cache(self, cache_name, method(self, *args, **kwargs))
- return get_from_cache(self, cache_name)
+ kwargs_pkl = pickle.dumps(kwargs)
+ if not _is_in_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl):
+ _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
+ return _get_from_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl)
return g
-
-
-def is_cached(self, name):
- """
- Determine if a cached item has been computed
- """
- return hasattr(self, "_memoize_cache") and name in self._memoize_cache.keys()
