CODEWITHSUNDEEP
pythonclasspython-3.4memoization
user
user

Reputation: 5696

How to memoize method when one of its memoized parameters should be an instance variable

To memoize f(x) I can use functools.lru_cache(): 

class A(object):

    def __init(self):
        self.time = 10  # This changes in the various spots in the program

    @functools.lru_cache(maxsize=None)
    def f(self, x):
        # Lots of code
        # ...
        # ...
        return x * some_other_func(self.time)

(to my understanding) lru_cache() creates a dict with various x as keys and their corresponding f(x) as values, so that it returns the stored values if I call f() with the same argument value, instead of recalculating it. However, this is not what I need.

My goal is to memoize f() values for different values of both x and self.time.

Using the following code achieves my goal: 

class A(object):

    def __init(self):
        self.time = 10

    @functools.lru_cache(maxsize=None)
    def g(self, x, t):
        # Lots of code
        # ...
        # ...
        return x * some_other_func(self.time)

    def f(self, x):
        return self.g(x=x, t=self.time)

Now instead of memoizing f(x) directly, I memoize g() which is always called with t=self.time.

However I am not sure if this is the cleanest solution. I would expect to only use a decorator for memoization and not have to create intermediate methods.

Is there a less messy way to achieve the above? (I will have to do the above for several methods, so I am looking for a solution as clean as possible)

Upvotes: 1

Views: 1010

Answers (1)

jonrsharpe
jonrsharpe

Reputation: 122061

If you want to implement this as a method decorator, here's one option that takes the state-relevant attribute names as arguments:

from functools import wraps

def stateful_memoize(*attrs):
    """Memoization that respects specified instance state."""
    def decorator(method):
        @wraps(method)
        def wrapper(self, *args):
            state = args + tuple(getattr(self, attr)
                                 for attr in attrs)
            if state not in wrapper.cache:
                wrapper.cache[state] = method(self, *args)
            return wrapper.cache[state]
        wrapper.cache = {}
        return wrapper
    return decorator

This simple version won't work with keyword method arguments, but otherwise should be fine. In use:

>>> class A(object):

    def __init__(self):
        self.time = 10

    @stateful_memoize('time')
    def f(self, x):
        print('Calling f with x={!r}, self.time={!r}'.format(x, self.time))
        return x * self.time

>>> a = A()
>>> a.f(1)
Calling f with x=1, self.time=10
10
>>> a.f(1)
10
>>> a.time = 5
>>> a.f(1)
Calling f with x=1, self.time=5
5
>>> a.time = 10
>>> a.f(1)
10
>>> a.f(2)
Calling f with x=2, self.time=10
20
>>> a.f.cache
{(1, 10): 10, (1, 5): 5, (2, 10): 20}

Upvotes: 2

Related Questions