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.

Toting

you can use class decorator if you really want to use the class based decorator but that will be again a separate independent class. – Scathe

Answer 1

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

Puttier

@Fleam I sure hope that's not the reason. I hardly care about the downvotes, but... I hope people realize they can just move the decorator into the class and it'll work just as well... – Puttier

Yes. The answer is fine, actually it's the same as my recommendation (see the last paragraph of my answer). – Fleam

I may have been a bit unclear when I asked my question, hence the edit. This is a great solution. – Toting

Unfortunately this solution does not preserve function signature with @functools.wraps(func) – Gytle

Don't declare self outside of the class. It doesn't makes sense to have a 'function' (which is not a 'method' inside a class), whose first argument is 'self'. Besides, if the decorator is used over a classmethod, then the self parameter would have a cls reference. So even though your solution may work, but it is syntactically wrong. I think this is why your answer has been getting down-votes. @Puttier CC: @Toting – Emetic

@ArchitKapoor You are right that self can be misleading here. The first parameter could be renamed (for example) obj and the code would keep the same functionality. Although I still cannot find a way to do it that generates 0 warnings, the trick proposed here does the job. Feel free to post a new Python-3 compatible answer if you know a more elegant way, as I believe myself and other people can still benefit from it. Thanks! – Toting

Answer 2

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.

Fleam

Agreed. If left in the class definition, I'd a least indicate it was private by changing its name to _enter_exit_info. – Bacchanalia

Answer 3

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

Papilionaceous

Answer 4

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

Semitics

Sorry dude, but this is not the point of my question. Besides, the answer is poorly edited. – Toting