pythondecoratorpython-2.6memoizationpython-decorators

Conditional memoization?


There are plenty of memoization decorators out there, but I'm curious how one would write a memoization decorator which supports arbitrary function signatures, but lets the function decide when to memoize a result? Something like this:

def conditional_memoize(f):
    cache = {}
    @wraps(f)
    def conditional_f(*args, **kwargs):
        return f(*args, **kwargs)
    return conditional_f

@conditional_memoize
def my_func(a, b, c):
    if str(a) + str(b) + str(c) in cache:
        return cache[str(a) + str(b) + str(c)]
    res = # compute the result
    if some_arbitrary_condition:
        cache[str(a) + str(b) + str(c)] = res
    return res

However, I know this won't work because of the NameError. Is there a clever approach to the problem anyway? I could always use a class method and a class cache, just wanted to see if there was a decorator pattern for this.


Solution

  • Have the function return both its desired result and a flag indicating whether the result should be cached, or have the wrapper pass the cache object to the function. (Or both!) Either way would work, but I like the first approach better. Maybe something like this...

    import functools
    
    def conditional_memoize(fn):
        cache = {}
    
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            key = args + tuple(sorted(kwargs.iteritems()))
            if key in cache:
                return cache[key]
            result, flag = fn(*args, **kwargs)
            if flag:
                cache[key] = result
            return result
    
        return wrapper