different behavior when implementing a decorator in Python with a function or a class

Question

I would like to write a decorator to be applied on methods of a class. The decorator should maintain a state therefore I would like to implement it with a class. However, when there are nested calls the class decorator fails while the decorator build with a function works.

here is a simple example :

def decorator(method):
    def inner(ref, *args, **kwargs):
        print(f'do something with {method.__name__} from class {ref.__class__}')
        return method(ref, *args, **kwargs)

    return inner


class class_decorator:

    def __init__(self, method):
        self.method = method

    def __call__(self, *args, **kwargs):
        print('before')
        result = self.method(*args, **kwargs)
        print('after')
        return result


class test:

    #@decorator
    @class_decorator
    def pip(self, a):
        return a + 1

    #@decorator
    @class_decorator
    def pop(self, a):
        result = a + self.pip(a)
        return result

t = test()
    
print(f'result pip : {t.pip(3)}')
print(f'result pop : {t.pop(3)}')

This will work with the 'decorator' function but not with the class_decorator because the nest call in the 'pop' method

Answer 1

The problem you are facing is because decorators of class methods are not passed methods, but functions.

In Python a method and a functions are two distinct types:

Python 3.8.3 (default, May 17 2020, 18:15:42)
Type 'copyright', 'credits' or 'license' for more information
IPython 7.15.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: class X:
   ...:     def m(self, *args, **kwargs):
   ...:         return [self, args, kwargs]

In [2]: type(X.m)
Out[2]: function

In [3]: type(X().m)
Out[3]: method

In [4]: X.m(1,2,x=3)
Out[4]: [1, (2,), {'x': 3}]

In [5]: X().m(1,2,x=3)
Out[5]: [<__main__.X at 0x7f1424f33a00>, (1, 2), {'x': 3}]

The "magic" transformation from a function (as m is in X) to a method (what it becomes when looked up in an instance X()) happens when m is looked up in the instance. Not being found in the instance itself Python looks it up in the class but when it discovers it is a function the value returned to who requested X().m is "wrapped up" in a method object that incorporates the self value.

The problem you are facing is however that this magic transformation is only applied if the value looked up ends up being a function object. If it's the instance of a class implementing __call__ (like in your case) the wrapping does not happen and thus the self value needed is not bound and the final code does not work.

A decorator should always return a function object and not a class instance pretending to be a function. Note that you can have all the state you want in a decorator because function objects in Python are actually "closures" and they can capture mutable state. For example:

In [1]: def deco(f):
   ...:     state = [0]
   ...:     def decorated(*args, **kwargs):
   ...:         state[0] += 1
   ...:         print(state[0], ": decorated called with", args, **kwargs)
   ...:         res = f(*args, **kwargs)
   ...:         print("return value", res)
   ...:         return res
   ...:     return decorated

In [2]: class X:
   ...:     def __init__(self, x):
   ...:         self.x = x
   ...:     @deco
   ...:     def a(self):
   ...:         return self.x + 1
   ...:     @deco
   ...:     def b(self):
   ...:         return 10 + self.a()

In [3]: x = X(12)

In [4]: x.a()
1 : decorated called with (<__main__.X object at 0x7f30a76f41c0>,)
return value 13
Out[4]: 13

In [5]: x.a()
2 : decorated called with (<__main__.X object at 0x7f30a76f41c0>,)
return value 13
Out[5]: 13

In [6]: x.b()
1 : decorated called with (<__main__.X object at 0x7f30a76f41c0>,)
3 : decorated called with (<__main__.X object at 0x7f30a76f41c0>,)
return value 13
return value 23
Out[6]: 23

In the above I used a simple list state but you can use as much state as you want including class instances. The important point is however that what decorators return is a function object. This way when looked up in a class instance the Python runtime will build the proper method object to make method calls work.

Another very important point to consider is however that decorators are executed at class definition time (i.e. when the class object is built) and not at instance creation. This means that the state that you will have in the decorator will be shared between all instances of the class.

Also another fact that may be not obvious and has bitten me in the past is that special methods like __call__ or __add__ are NOT looked up in the instance first and Python goes directly in the class object to look them up. This is a documented implementation choice but is none the less a "strange" asymmetry that can come as a surprise.

Answer 2

Decorator is just "syntactic sugar". The problem with the class decorator, is that self is no longer passed as the first argument.

What we want, is to mimic the behavior of decorator, in which we return a methd that no longer needs self to be passed to it.

This is something that can be done directly with the partial function, by making it a descriptor

You will notice that the first function called is the __get__

class class_decorator:

    def __init__(self, method):
        self.method = method
        
    def __set_name__(self, owner, name):
        self.owner = owner

    def __call__(self, *args, **kwargs):
        print('before')
        result = self.method(*args,**kwargs)
        print('after')
        return result
        
    def __get__(self, instance, owner):
        print('calling get')
        from functools import partial
        return partial(self, instance)