Search code examples
pythonpython-2.7wrapperdecoratorpython-decorators

Declaring decorator inside a class


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.


Solution

  • 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()