Summary: In this tutorial, we will learn to decorate a class in Python using the function and class decorator.

A decorator in Python is any callable object that modifies existing functions or classes.

It modifies the old object with additional Python statements and returns the reference of the same.

For example, consider the following class, it has two methods: __init__ and display. The __init__ method initializes the name variable while the display outputs the name:

class Student:
    def __init__(self, name):
        self.name = name
        
    def display(self):
        print('Name:', self.name)

To decorate this class in Python, we can either add new methods to the class or modify the existing ones, or do both.

Also, there are two ways of doing so in Python, either by using a function decorator or by using a class decorator.

Let’s see an example of each one by one.

Decorate Class using Function Decorator

To decorate the class using a function decorator accept the class as an argument, modify its code and return the class at the end.

def mydecorator(student):
    #define a new display method
    def newdisplay(self):
        print('Name: ', self.name)
        print('Subject: Programming')
    
    #replace the display with newdisplay
    student.display = newdisplay
    
    #return the modified student 
    return student

@mydecorator
class Student:
    def __init__(self, name):
        self.name = name
        
    def display(self):
        print('Name:', self.name)
        

obj = Student('Pencil Programmer')
obj.display()
'''
Name:  Pencil Programmer
Subject: Programming
'''

If the display method did not existed in the class, the newdisplay would have been added to the class as the display method.

Since the reference of the class is available in the decorator function, we can add new attributes and methods to the class in addition of modifying the existing methods.

Decorate Class using Class Decorator

To decorate a class using a class decorator accept the reference of the class as an argument (in the __init__ method of the decorator), modify its code in the __call__ method, and finally return the instance of the modified class.

class Mydecorator:
    #accept the class as argument
    def __init__(self, student):
        self.student = student
    
    #accept the class's __init__ method arguments
    def __call__(self, name):
        #define a new display method
        def newdisplay(self):
            print('Name: ', self.name)
            print('Subject: Python')
            
        #replace display with newdisplay
        self.student.display = newdisplay
        
        #return the instance of the class
        obj = self.student(name)
        return obj

@Mydecorator
class Student:
    def __init__(self, name):
        self.name = name
        
    def display(self):
        print('Name: ', self.name)

obj = Student('Pencil Programmer')
obj.display()
'''
Name:  Pencil Programmer
Subject: Python
'''

The only difference here is that instead of class reference we are returning the reference of the object.

Leave a Reply