python memoization decorator

memoize.py

'''
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