add a decorate function to a class

Question

I have a decorated function (simplified version):

class Memoize:
    def __init__(self, function):
        self.function = function
        self.memoized = {}
    def __call__(self, *args, **kwds):
        hash = args
        try:
            return self.memoized[hash]
        except KeyError:
            self.memoized[hash] = self.function(*args)
            return self.memoized[hash]


@Memoize
def _DrawPlot(self, options):
    do something...

now I want to add this method to a pre-esisting class.

ROOT.TChain.DrawPlot = _DrawPlot

when I call this method:

chain = TChain()
chain.DrawPlot(opts)

I got:

self.memoized[hash] = self.function(*args)
TypeError: _DrawPlot() takes exactly 2 arguments (1 given)

why doesn't it propagate self?

Answer 1

The problem is that you have defined your own callable class then tried to use it as a method. When you use a function as an attribute, accessing the function as an attribute calls it its __get__ method to return something other than the function itself—the bound method. When you have your own class without defining __get__, it just returns your instance without implicitly passing self.

Descriptors are explained some on http://docs.python.org/reference/datamodel.html#descriptors if you are not familiar with them. The __get__, __set__, and __delete__ methods change how interacting with your object as an attribute works.

---

You could implement memoize as a function and use the built-in __get__ magic that functions already have

import functools

def memoize(f):
    @functools.wraps(f)
    def memoized(*args, _cache={}): 
        # This abuses the normally-unwanted behaviour of mutable default arguments.
        if args not in _cache:
            _cache[args] = f(*args)
        return _cache[args]
    return memoized

or by modifying your class along the lines of

import functools

class Memoize(object): #inherit object
    def __init__(self, function):
        self.function = function
        self.memoized = {}
    def __call__(self, *args): #don't accept kwargs you don't want.
        # I removed "hash = args" because it shadowed a builtin function and 
        # because it was untrue--it wasn't a hash, it was something you intended for
        # Python to hash for you.
        try:
            return self.memoized[args]
        except KeyError:
            self.memoized[args] = self.function(*args)
            return self.memoized[args]
    def __get__(self, obj, type):
        if obj is None: #We looked up on the class
            return self

        return functools.partial(self, obj)

Note that both of these choke if any of the arguments you pass in are mutable (well, unhashable technically). This might be suitable for your case, but you may also want to deal with the case where args is unhashable.