Conditional memoization?

Question

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.

Answer 1

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