Python - Class Decorators

We have used functions to decorate functions and to decorate classes. Now, we will see how to define a class as a decorator.

At the start of this chapter, in the definition of decorator, we had seen that a decorator is a callable; a callable is any object that can be called a function. If a class has the __call__ method defined inside it, then instance objects of that class become callable objects. When these instance objects are called like a function using parentheses, the code inside the __call__ method is executed. Now, let us see how we can use such a class as a decorator.

class MyDecorator:
    def __init__(self, fn):
        self.fn = fn

    def __call__(self,*args, **kwargs):
        print('Decoration before execution of function') 
        self.fn(*args, **kwargs)
        print('Decoration after execution of function\n') 


def func(message, name):
    print(message, name)

func = MyDecorator(func)

The class MyDecorator is going to serve as a decorator. We have a simple function named func and we have decorated it using the class MyDecorator by using the manual decoration syntax. In the statement func = MyDecorator(func), we are sending the undecorated function to the class and getting back the decorated function. We know that when we call a class object ( MyDecorator(func) ), the __init__ method of the class is executed. This is why we have made __init__ method accept a function. The call MyDecorator(func) returns an instance object of the class. So, after the statement func = MyDecorator(func), func is actually an instance object of the MyDecorator class.

So, this is how we can create a decorator class. The second parameter to __init__ accepts the function to be decorated, and this function is stored as an instance variable. Inside the __call__ method the original undecorated function is called and the decoration code is also placed before or after the call. The class is instantiated at the decoration time, and the instance object that is created is assigned to the function name. After decoration, the function name refers to an instance object, and so when the function name is called, the __call__ method of the class is executed.

Instead of using the manual decoration statement, we can use the automatic decoration syntax to get the same effect.

@MyDecorator
def func(message, name):
    print(message, name)

 

You can return the result from the __call__ method, so that the return value of the original function is not lost.

def __call__(self,*args, **kwargs):
    print('Decoration before execution of function') 
    result = self.fn(*args, **kwargs)
    print('Decoration after execution of function\n') 

    return result

If you want to preserve the metadata, then you have to call the wraps decorator.

def __init__(self, fn):
    self.fn = fn
    wraps(fn)(self)
    self.author = 'Jim'   # Add a new attribute here

If you want the decorator to add an attribute to the function, you have to add it inside the __init__ method, not inside the __call__ method. This is because we do not want to add the attribute each time the function is called. We want to add the attribute just once when the function is decorated.