Search code examples
pythonclassmethodsdecorator

Using class as decorator for another class's method


I have an issue with using a class to decorate another class' method. Code is as follows:

class decorator(object):
    def __init__(self, func):
        self.func = func

    def __call__(self, *args):
        return self.func(*args)

class test(object):
    @decorator
    def func(self, x, y):
        print x, y

t = test()
t.func(1, 2)

It shows this error

TypeError: func() takes exactly 3 arguments (2 given).

If called using:

t.func(t, 1, 2)

then it passes. But then if the decorator is taken away, then this line will have issue again.

Why this is happening and how to solve it?

Edit: second version of the code to show the self in decorator.__call__ should be different than the self in test.func:

class decorator(object):
    def __init__(self, func):
        self.func = func

    def __call__(self, *args):
        return self.func(*args)

class test(object):
    def __init__(self):
        self.x = 1
        self.y = 2
    @decorator
    def func(self):
        print self
        print self.x, self.y

t = test()
t.func()

This shows the same error. But

t.func(t)

works but not ideal.


Solution

  • To work as a method, an object in a class needs to implement part of the descriptor protocol. That is, it should have a __get__ method that returns a callable object which has been "bound" to the instance the method was looked up on.

    Here's one way that you could make that work, using a wrapper function:

    class decorator(object):
        def __init__(self, func):
            self.func = func
    
        def __get__(self, instance, owner):
            def wrapper(*args):
                return self.func(instance, *args) # note, self here is the descriptor object
            return wrapper
    

    You could instead return an instance of some other class from __get__, rather than a function, and use the __call__ method of that other class to implement the wrapper. If you're not using a closure though, you'd need to pass the instance to the wrapper class explicitly (as well as the function, since self.func won't work outside the descriptor class).