''' Example: # Memoize result from `myfunction1` depending on the arg1 and arg3 parameters. @memoize('arg1', 'arg3') def myfunction1(arg1, arg2, arg3): ... res = myfunction1('key1', 'meh', 7) # Memoize result from `myfunction2` depending on 'prop' attribute from "argobject" parameter. class myobject(object): prop = 'hi' @memoize(argobject='prop') def myfunction2(blah, argobject): .... x = myobject() x.prop = 'heh' res = myfunction2('whee', x) # Memoize result from `myfunction3` depending on sum of the "list_of_numbers" parameter @memoize(list_of_numbers=lambda list: sum(list)) def myfunction3(blargh, list_of_numbers): ... res = myfunction3(None, [1,3,5,7]) ''' ######################################################################## # memoization decorator for arbitrary parameters # ######################################################################## def memoize(*kargs, **kattrs): '''Converts a function into a memoized callable kargs = a list of positional arguments to use as a key kattrs = a keyword-value pair describing attributes to use as a key if key='string', use kattrs[key].string as a key if key=callable(n)', pass kattrs[key] to callable, and use the returned value as key if no memoize arguments were provided, try keying the function's result by _all_ of it's arguments. ''' F_VARARG = 0x4 F_VARKWD = 0x8 F_VARGEN = 0x20 kargs = list(kargs) kattrs = tuple((o, a) for o, a in sorted(kattrs.items())) # Define some utility functions for interacting with a function object in a portable manner has_function = (lambda F: hasattr(F, 'im_func')) if sys.version_info.major < 3 else (lambda F: hasattr(F, '__func__')) get_function = (lambda F: F.im_func) if sys.version_info.major < 3 else (lambda F: F.__func__) def prepare_callable(fn, kargs=kargs, kattrs=kattrs): if has_function(fn): fn = get_function(fn) if not isinstance(fn, memoize.__class__): raise AssertionError("Callable {!r} is not of a function type".format(fn)) cache = {} co = fn.__code__ flags, varnames = co.co_flags, iter(co.co_varnames) if (flags & F_VARGEN) != 0: raise AssertionEerror("Not able to memoize {!r} generator function".format(fn)) argnames = itertools.islice(varnames, co.co_argcount) c_positional = tuple(argnames) c_attribute = kattrs c_var = (six.next(varnames) if flags & F_VARARG else None, six.next(varnames) if flags & F_VARKWD else None) if not kargs and not kattrs: kargs[:] = itertools.chain(c_positional, filter(None, c_var)) def key(*args, **kwds): res = iter(args) p = dict(zip(c_positional, res)) p.update(kwds) a, k = c_var if a is not None: p[a] = tuple(res) if k is not None: p[k] = dict(kwds) k1 = (p.get(k, None) for k in kargs) k2 = ((n(p[o]) if callable(n) else getattr(p[o], n, None)) for o, n in c_attribute) return tuple(itertools.chain(k1, [None], k2)) def callee(*args, **kwds): res = key(*args, **kwds) try: return cache[res] if res in cache else cache.setdefault(res, fn(*args, **kwds)) except TypeError: return fn(*args, **kwds) def force(*args, **kwds): res = key(*args, **kwds) cache[res] = fn(*args, **kwds) return cache[res] # set some utilies on the memoized function callee.memoize_key = lambda *args, **kwargs: key(*args, **kwargs) callee.memoize_key.__doc__ = """Generate a unique key based on the provided arguments.""" callee.memoize_cache = lambda: cache callee.memoize_cache.__doc__ = """Return the current memoize cache.""" callee.memoize_clear = lambda: cache.clear() callee.memoize_clear.__doc__ = """Empty the current memoize cache.""" callee.force = lambda *args, **kwargs: force(*args, **kwargs) callee.force.__doc__ = """Force calling the function whilst updating the memoize cache.""" callee.__name__ = fn.__name__ callee.__doc__ = fn.__doc__ callee.callable = fn return callee if isinstance(callee, types.FunctionType) else types.FunctionType(callee) return prepare_callable(kargs.pop(0)) if not kattrs and len(kargs) == 1 and callable(kargs[0]) else prepare_callable