The Kitchin Research Group

Suppose you have a module that you import a class from, and the class defines some methods that you want to memoize. You do not want to modify the source code, maybe because it is not your code, or because you do not want to maintain it, etc… Here is one way to modify the class functions at runtime. We will use the memoize decorator and replace the class function definition with the wrapped function that caches the results. We also allow arbitrary arguments and keyword arguments. A subtle wrinkle here is that you cannot use a dictionary as a key to a dictionary because dictionaries are not hashable. We use the pickle module to created a string that should uniquely represent the args and keyword args, and we use that string as the key.

class Calculator:
    def __init__(self):
        pass

    def calculate(self, a):
        'returns the answer to everything'
        return 42

    def method_2(self, *args, **kwargs):
        return (args, kwargs)


import pickle

from functools import wraps
def memoize(func):
    cache = {}
    @wraps(func)
    def wrap(*args,**kwargs):
        key = pickle.dumps((args, kwargs))
        if key not in cache:
            print 'Running func with ', args
            cache[key] = func(*args, **kwargs)
        else:
            print 'result in cache'
        return cache[key]
    return wrap

# now monkey patch/decorate the class function
Calculator.calculate = memoize(Calculator.calculate)
Calculator.method_2 = memoize(Calculator.method_2)

calc = Calculator()
print calc.calculate(3)
print calc.calculate(3)
print calc.calculate(4)
print calc.calculate(3)


print calc.method_2()
print calc.method_2()

print calc.method_2(1,2)
print calc.method_2(1,2)

print calc.method_2(1,2,a=5)
print calc.method_2(1,2,a=5)

Running func with  (<__main__.Calculator instance at 0x0000000001E9B3C8>, 3)
42
result in cache
42
Running func with  (<__main__.Calculator instance at 0x0000000001E9B3C8>, 4)
42
result in cache
42
Running func with  (<__main__.Calculator instance at 0x0000000001E9B3C8>,)
((), {})
result in cache
((), {})
Running func with  (<__main__.Calculator instance at 0x0000000001E9B3C8>, 1, 2)
((1, 2), {})
result in cache
((1, 2), {})
Running func with  (<__main__.Calculator instance at 0x0000000001E9B3C8>, 1, 2)
((1, 2), {'a': 5})
result in cache
((1, 2), {'a': 5})

This particular memoize decorator is not persistent; the data is only stored in memory. You would have to write the data out to a file and reread the file to make it persistent.

It is not obvious this practice is good; you have in essence changed the behavior of the original function in a way that may be hard to debug, and could conceivably be incompatible with the documentation of the function.

An alternative approach is writing another function that wraps the code you want, and memoize that function.

class Calculator:
    def __init__(self):
        pass

    def calculate(self, a):
        'returns the answer to everything'
        return 42



from functools import wraps
def memoize(func):
    cache = {}
    @wraps(func)
    def wrap(*args):
        if args not in cache:
            print 'Running func with ', args
            cache[args] = func(*args)
        else:
            print 'result in cache'
        return cache[args]
    return wrap

calc = Calculator()

@memoize
def my_calculate(a):
    return calc.calculate(a)

print my_calculate(3)
print my_calculate(3)
print my_calculate(4)
print my_calculate(3)

Running func with  (3,)
42
result in cache
42
Running func with  (4,)
42
result in cache
42

It is debatable whether this is cleaner. One argument for this is that it does not monkey with the original code at all.

Copyright (C) 2013 by John Kitchin. See the License for information about copying.

org-mode source