Python: How to Use Class Decorators with Dataclass

This concise, straight-to-the-point article shows you how to use class decorators with dataclass in Python. We’ll take a look at the fundamentals and then walk through a couple of practical examples.

The Fundamentals

Decorators can be used with dataclasses in Python to modify the behavior of the dataclass or add new functionality to it. Here are some scenarios where you might want to use decorators with dataclasses:

• Adding properties or methods: You can use class decorators to add new properties or methods to a dataclass. This can be useful if you want to add functionality to a dataclass without modifying its source code.
• Changing behavior: You can use class decorators to change the behavior of a dataclass method. For example, you might want to change the way a post-init method works for a particular dataclass.
• Specifying types of attributes: You can use descriptors with class decorators to specify the type of an attribute in a dataclass. This can help ensure that the values assigned to the attribute are of the correct type, which can be particularly useful in large or complex dataclasses.

You can name the decorators however you want. The name you give to a decorator is simply the name of the function that defines the decorator. When you use the decorator syntax ( @decorator_name ), you’re calling the function that defines the decorator and passing the decorated object (a function or a class) as an argument.

Words might be confusing and hard to follow at first; let’s write some code for a better understanding.

Examples

Adding a property to a dataclass

The code:

from dataclasses import dataclass

# Decorator to add a property to our dataclass
# This is a class decorator, so it takes a class as an argument
def add_property(cls):
    cls.new_property = property(lambda self: self.some_property * 4)
    return cls

@add_property 
@dataclass
class MyClass:
    some_property: int

my_instance = MyClass(some_property=5)
print(my_instance.new_property)

Output:

20

Adding a method to a dataclass

The code:

from dataclasses import dataclass

# Create a decorator that adds a method to a class
# The decorator takes a class as an argument
def add_method(cls):
    def new_method(self):
        return self.some_property ** 2
    cls.new_method = new_method
    return cls

# Use the decorator to add a method to our dataclass
@add_method
@dataclass
class MyClass:
    some_property: int

my_instance = MyClass(some_property=9)
result = my_instance.new_method()
print(result)

Output:

81

Changing the behavior of a dataclass method

The code:

from dataclasses import dataclass

# Define a function that doubles the age of a person
def double_age(person):
    person.age *= 2
    return person

# Define a function that modifies the post_init method
def post_init(self):
    self = double_age(self)

# Define a decorator that modifies the post_init method
def modify_post_init(cls):
    original_post_init = cls.__init__
    def new_post_init(self, *args, **kwargs):
        original_post_init(self, *args, **kwargs)
        post_init(self)
    cls.__init__ = new_post_init
    return cls

# Define a dataclass that uses the decorator
@modify_post_init
@dataclass
class Person:
    name: str
    age: int

person = Person("Pipi", 35)
print(person.age)

Output:

70

Final Words

Class decorators can be a powerful tool when working with dataclasses in Python. They allow you to modify the behavior of a class without having to modify its source code. If they are used properly, your programs will be more flexible, easier to maintain, and better organized.