How to use joblib.Memory of cache the output of a member function of a Python Class

Question

I would like to cache the output of a member function of a class using joblib.Memory library. Here is a sample code:

import joblib
import numpy as np

mem = joblib.Memory(cachedir='/tmp', verbose=1)


@mem.cache
def my_sum(x):
    return np.sum(x)


class TestClass(object):
    def __init__(self):
        pass

    @mem.cache
    def my_sum(self, x):
        return np.sum(x)


if __name__ == '__main__':
    x = np.array([1, 2, 3, 4])
    a = TestClass()
    print a.my_sum(x)  # does not work
    print my_sum(x) # works fine

However, I get the following error:

/nfs/sw/anaconda2/lib/python2.7/site-packages/joblib/memory.pyc in _get_output_dir(self, *args, **kwargs)
    512             of the function called with the given arguments.
    513         """
--> 514         argument_hash = self._get_argument_hash(*args, **kwargs)
    515         output_dir = os.path.join(self._get_func_dir(self.func),
    516                                   argument_hash)

/nfs/sw/anaconda2/lib/python2.7/site-packages/joblib/memory.pyc in _get_argument_hash(self, *args, **kwargs)
    505     def _get_argument_hash(self, *args, **kwargs):
    506         return hashing.hash(filter_args(self.func, self.ignore,
--> 507                                          args, kwargs),
    508                              coerce_mmap=(self.mmap_mode is not None))
    509 

/nfs/sw/anaconda2/lib/python2.7/site-packages/joblib/func_inspect.pyc in filter_args(func, ignore_lst, args, kwargs)
    228                            repr(args)[1:-1],
    229                            ', '.join('%s=%s' % (k, v)
--> 230                                     for k, v in kwargs.items())
    231                            )
    232                         )

ValueError: Wrong number of arguments for my_sum(self, x):
     my_sum(array([1, 2, 3, 4]), ) was called.

Is there a way to cache a member function of a class using Memory or any other decorators?

Answer 1

The following excerpt is taken from https://joblib.readthedocs.io/en/latest/memory.html#gotchas

caching methods: you cannot decorate a method at class definition, because when the class is instantiated, the first argument (self) is bound, and no longer accessible to the Memory object. The following code won’t work:

class Foo(object):

    @mem.cache  # WRONG
    def method(self, args):
        pass

The right way to do this is to decorate at instantiation time:

class Foo(object):

    def __init__(self, args):
        self.method = mem.cache(self.method)

    def method(self, ...):
        pass