Last active
August 29, 2015 14:07
-
-
Save chronossc/c940f81bc78fd55708c4 to your computer and use it in GitHub Desktop.
memoize.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# -*- coding: utf-8 -*- | |
from __future__ import (absolute_import, division, print_function, | |
unicode_literals) | |
from functools import partial, update_wrapper | |
try: | |
from cPickle import dumps | |
except ImportError: | |
from pickle import dumps | |
class DictCache(dict): | |
def set(self, key, value, *args, **kwargs): | |
""" | |
Set value to dict ignoring args and kwargs that dict will not use. | |
>>> d = DictCache() | |
>>> d.set('foo', 'bar') | |
>>> assert d['foo'] == 'bar' | |
""" | |
self[key] = value | |
def delete(self, key): | |
""" | |
Delete a item from dict. | |
>>> d = DictCache() | |
>>> d.set('foo', 'bar') | |
>>> assert d['foo'] == 'bar' | |
>>> d.delete('foo') | |
>>> assert 'foo' not in d | |
""" | |
if key in self: | |
del self[key] | |
_keymaker = lambda *a, **kw: dumps((a, kw)) | |
class Memoize(object): | |
""" | |
Decorator that caches a function's return value each time it is called. | |
If called later with the same arguments, the cached value is returned | |
(not reevaluated). For backends that not implement __contains__ we just call | |
cache.get and ignore result if it is None. | |
param func: the function to be memoized | |
param cache: Callable returning the cache storage to be used. Storage should support get, set | |
and delete methods. Default is a dictionary cache. It's lazy. | |
param keymaker: Callable that receive func args and kwargs to create a key. | |
params **kwargs: Parameters to be passed to cache.set methods. | |
""" | |
def __init__(self, func, cache=None, keymaker=None, **kwargs): | |
self._func = func | |
self._cache = None | |
self._keymaker = keymaker or _keymaker # use pickle keymaker if no one is passed | |
self._cache_callable = cache or DictCache | |
self._set_kwargs = kwargs | |
update_wrapper(self, func) | |
@property | |
def cache(self): | |
# make cache lazy | |
if self._cache is None: | |
self._cache = self._cache_callable() | |
return self._cache | |
def __call__(self, *args, **kwargs): | |
key = self._keymaker(*args, **kwargs) | |
try: | |
# For iterable backends, we use 'in' (which call __contains__). | |
if key in self.cache: | |
return self.cache.get(key) | |
except TypeError: | |
# For non iterable backends we call get. cache.get normally returns | |
# None if not found, so for non iterable backends we don't memoize | |
# results that is None. | |
cached = self.cache.get(key) | |
if cached: | |
return cached | |
value = self._func(*args, **kwargs) | |
self.cache.set(key, value, **self._set_kwargs) | |
return value | |
def __repr__(self): | |
""" Return the function's repr """ | |
return self._func.__repr__() | |
def __get__(self, obj, objtype): | |
""" Support instance methods """ | |
return partial(self.__call__, obj) | |
def memoize(function=None, cache=None, keymaker=None, **kwargs): | |
""" | |
>>> @memoize | |
... def fibonacci(n): | |
... "Return the nth fibonacci number." | |
... if n in (0, 1): | |
... return n | |
... return fibonacci(n - 1) + fibonacci(n - 2) | |
>>> print(fibonacci(12)) | |
144 | |
>>> calls = 0 | |
>>> @memoize(keymaker=lambda a, b: '_'.join([a, b])) | |
... def dosomething(a, b): | |
... global calls | |
... calls += 1 | |
... return a, b | |
>>> dosomething(u'1', u'2') | |
(u'1', u'2') | |
>>> assert calls == 1 | |
>>> dosomething('1', '2') | |
(u'1', u'2') | |
>>> assert calls == 1 | |
""" | |
if function: | |
return Memoize(function) | |
else: | |
def wrapper(function): | |
return Memoize(function, cache=cache, keymaker=keymaker, **kwargs) | |
return wrapper |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment