Skip to content

Instantly share code, notes, and snippets.

@chronossc
Last active August 29, 2015 14:07
Show Gist options
  • Save chronossc/c940f81bc78fd55708c4 to your computer and use it in GitHub Desktop.
Save chronossc/c940f81bc78fd55708c4 to your computer and use it in GitHub Desktop.
memoize.py
# -*- 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