Declaring decorator inside a class

Question

I'm trying to use custom wrappers/decorators in Python, and I'd like to declare one inside a class, so that I could for instance print a snapshot of the attributes. I've tried things from this question with no success.

---

Here is what I'd like to do (NB: this code doesn't work, I explain what happens below)

class TestWrapper():
    def __init__(self, a, b):
        self.a = a
        self.b = b
        self.c = 0

    def enter_exit_info(self, func):
        def wrapper(*arg, **kw):
            print '-- entering', func.__name__
            print '-- ', self.__dict__
            res = func(*arg, **kw)
            print '-- exiting', func.__name__
            print '-- ', self.__dict__
            return res
        return wrapper

    @enter_exit_info
    def add_in_c(self):
        self.c = self.a + self.b
        print self.c

    @enter_exit_info
    def mult_in_c(self):
        self.c = self.a * self.b
        print self.c


if __name__ == '__main__':
    t = TestWrapper(2, 3)
    t.add_in_c()
    t.mult_in_c()

The expected output is :

-- entering add_in_c
-- {'a': 2, 'b': 3, 'c': 0}
5
-- exiting add_in_c
-- {'a': 2, 'b': 3, 'c': 5}
-- entering mult_in_c
-- {'a': 2, 'b': 3, 'c': 5}
6
-- exiting mult_in_c
-- {'a': 2, 'b': 3, 'c': 6}

But I this code gives

Traceback (most recent call last):
  File "C:\Users\cccvag\workspace\Test\src\module2.py", line 2, in <module>
    class TestWrapper():
  File "C:\Users\cccvag\workspace\Test\src\module2.py", line 18, in     TestWrapper
    @enter_exit_info
TypeError: enter_exit_info() takes exactly 2 arguments (1 given)

And if I try @enter_exit_info(self) or @self.enter_exit_info, I get a NameError. What could I do?

---

EDIT:

I do not need above all to have the decorator physically declared inside the class, as long as it is able to access attributes from an instance of this class. I thought it could only be made by declaring it inside the class, Rawing's answer proved me wrong.

Answer 1

You will need to handle self explicitly.

class TestWrapper:
    def __init__(self, a, b):
        self.a = a
        self.b = b
        self.c = 0

    def enter_exit_info(func):
        def wrapper(self, *arg, **kw):
            print '-- entering', func.__name__
            print '-- ', self.__dict__
            res = func(self, *arg, **kw)
            print '-- exiting', func.__name__
            print '-- ', self.__dict__
            return res
        return wrapper

    @enter_exit_info
    def add_in_c(self):
        self.c = self.a + self.b
        print self.c

    @enter_exit_info
    def mult_in_c(self):
        self.c = self.a * self.b
        print self.c


if __name__ == '__main__':
    t = TestWrapper(2, 3)
    t.add_in_c()
    t.mult_in_c()

This is valid python, but it's somewhat weird to have a function at the class level which is not really a method. Unless you have a good reason to do it this way, it would be more idiomatic to move the decorator to module level scope.

Answer 2

Instead of defining the decorator inside the class you can just intercept the self parameter:

import functools

def enter_exit_info(func):
    @functools.wraps(func)
    def wrapper(self, *arg, **kw):
        print '-- entering', func.__name__
        print '-- ', self.__dict__
        res = func(self, *arg, **kw)
        print '-- exiting', func.__name__
        print '-- ', self.__dict__
        return res
    return wrapper

class TestWrapper():
    def __init__(self, a, b):
        self.a = a
        self.b = b
        self.c = 0
    
    @enter_exit_info
    def add_in_c(self):
        self.c = self.a + self.b
        print self.c

    @enter_exit_info
    def mult_in_c(self):
        self.c = self.a * self.b
        print self.c


if __name__ == '__main__':
    t = TestWrapper(2, 3)
    t.add_in_c()
    t.mult_in_c()

Answer 3

Hi do you want the output should be in dictionary format? If you don't want the output in dictionary format u can try this....

def enter_exit_info(func):
        def wrapper(*arg, **kw):
            print '-- entering', func.__name__        
            res = func(*arg, **kw)
            print '-- exiting', func.__name__
            return res
        return wrapper

then your output will be

-- entering add_in_c

5
-- exiting add_in_c

-- entering mult_in_c

6
-- exiting mult_in_c

Answer 4

TL;DR : what you want is

def enter_exit_info(func):
    def wrapper(self, *arg, **kw):
        print '-- entering', func.__name__
        print '-- ', self.__dict__
        res = func(*arg, **kw)
        print '-- exiting', func.__name__
        print '-- ', self.__dict__
        return res
    return wrapper

Remember that

@decorate
def myfunc():
    pass

is really just syntactic sugar for

def myfunc():
    pass
my_func = decorate(my_func)

So since in your case, decorated functions are replaced by the decorator's wrapper function, it's this wrapper function that will receive the current instance as first argument.

EDIT : I positively agree with other answers on the point that it makes no sense defining this decorator within the class. You don't need it to access the current instance since it's provided as the function's first argument. FWIW the def statement doesn't work any differently from being used within a class statement, it always yields a plain old function object. What makes the function a "method" (and 'automagically' pass the current instance as first argument) is the attribute resolution mechanism, cf https://wiki.python.org/moin/FromFunctionToMethod