Reduce boilerplate code with Python's dataclasses

Recently, my work needed me to create lots of custom data types and draw comparison among them. So, my code was littered with many classes that somewhat looked like this:

class CartesianPoint:
    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z

    def __repr__(self):
        return f"CartesianPoint(x = {self.x}, y = {self.y}, z = {self.z})"


print(CartesianPoint(1, 2, 3))

>>> CartesianPoint(x = 1, y = 2, z = 3)

This class only creates a CartesianPoint type and shows a pretty output of the instances created from it. However, it already has two methods inside, __init__ and __repr__ that do not do much.

Dataclasses

Let’s see how data classes can help to improve this situation. Data classes were introduced to python in version 3.7. Basically they can be regarded as code generators that reduce the amount of boilerplate you need to write while generating generic classes. Rewriting the above class using dataclass will look like this:

from dataclasses import dataclass


@dataclass
class CartesianPoint:
    x: float
    y: float
    z: float


# using the class
point = CartesianPoint(1, 2, 3)
print(point)

>>> CartesianPoint(x=1, y=2, z=3)

In the above code, the magic is done by the dataclass decorator. Data classes require you to use explicit type annotations¹ and it automatically implements methods like __init__, __repr__, __eq__ etc beforehand. You can inspect the methods that dataclass auto defines via Python’s help.

help(CartesianPoint)

Help on class CartesianPoint in module __main__:

class CartesianPoint(builtins.object)
 |  CartesianPoint(x:float, y:float, z:float)
 |
 |  Methods defined here:
 |
 |  __eq__(self, other)
 |
 |  __init__(self, x:float, y:float, z:float) -> None
 |
 |  __repr__(self)
 |
 |  ----------------------------------------------------------------------
 |  Data descriptors defined here:
 |
 |  __dict__
 |      dictionary for instance variables (if defined)
 |
 |  __weakref__
 |      list of weak references to the object (if defined)
 |
 |  ----------------------------------------------------------------------
 |  Data and other attributes defined here:
 |
 |  __annotations__ = {'x': <class 'float'>, 'y': <class 'float'>, 'z': <c...
 |
 |  __dataclass_fields__ = {'x': Field(name='x',type=<class 'float'>,defau...
 |
 |  __dataclass_params__ = _DataclassParams(init=True,repr=True,eq=True,or...
 |
 |  __hash__ = None

Using default values

You can provide default values to the fields in the following way:

from dataclasses import dataclass


@dataclass
class CartesianPoint:
    x: float = 0
    y: float = 0
    z: float = 0

Using arbitrary field type

If you don’t want to specify your field type during type hinting, you can use Any type from python’s typing module.

from dataclasses import dataclass
from typing import Any


@dataclass
class CartesianPoint:
    x: Any
    y: Any
    z: Any

Instance ordering

You can check if two instances are equal without making any modification to the class.

from dataclasses import dataclass


@dataclass
class CartesianPoint:
    x: float
    y: float
    z: float


point_1 = CartesianPoint(1, 2, 3)
point_2 = CartesianPoint(1, 2, 5)

print(point_1 == point_2)

>>> False

However, if you want to compare multiple instances of dataclasses, aka add __gt__ or __lt__ methods to your instances, you have to turn on the order flag manually.

from dataclasses import dataclass


@dataclass(order=True)
class CartesianPoint:
    x: float
    y: float
    z: float


# comparing two instances
point_1 = CartesianPoint(10, 12, 13)
point_2 = CartesianPoint(1, 2, 5)

print(point_1 > point_2)

>>> True

By default, while comparing instances, all of the fields are used. In our above case, all the fields x, y, zof point_1 instance are compared with all the fields of point_2 instance. You can customize this using the field function.

Suppose you want to acknowledge two instances as equal only when attribute x of both of them are equal. You can emulate this in the following way:

from dataclasses import dataclass, field


@dataclass(order=True)
class CartesianPoint:
    x: float
    y: float = field(compare=False)
    z: float = field(compare=False)


# create intance where only the x attributes are equal
point_1 = CartesianPoint(1, 3, 5)
point_2 = CartesianPoint(1, 4, 6)

# compare the instances
print(point_1 == point_2)
print(point_1 < point_2)

>>> True
>>> False

You can see the above code prints out True despite the instances have different y and z attributes.

Adding methods

Methods can be added to dataclasses just like normal classes. Let’s add another method called dist to our CartesianPoint class. This method calculates the distance of a point from origin.

from dataclasses import dataclass
import math


@dataclass
class CartesianPoint:
    x: float
    y: float
    z: float

    def dist(self):
        return math.sqrt(self.x**2 + self.y**2 + self.z**2)


# create a new instance and use method `abs_val`
point = CartesianPoint(5, 6, 7)
norm = point.abs_val()

print(norm)

>>> 10.488088481701515

Making instances immutable

By default, instances of dataclasses are immutable. If you want to prevent mutating your instance attributes, you can set frozen=True while defining your dataclass.

from dataclasses import dataclass


@dataclass(frozen=True)
class CartesianPoint:
    x: float
    y: float
    z: float

If you try to mutate the any of the attributes of the above class, it will raise FrozenInstanceError.

point = CartesianPoint(2, 4, 6)
point.x = 23

---------------------------------------------------------------------------

FrozenInstanceError                       Traceback (most recent call last)

<ipython-input-34-b712968bd0eb> in <module>
        1 point = CartesianPoint(2, 4, 6)
----> 2 point.x = 23


<string> in __setattr__(self, name, value)


FrozenInstanceError: cannot assign to field 'x'

Making instances hashable

You can turn on the unsafe_hash parameter of the dataclass decorator to make the class instances hashable. This may come in handy when you want to use your instances as dictionary keys or want to perform set operation on them. However, if you are using unsafe_hash make sure that your dataclasses do not contain any mutable data structure in it.

from dataclasses import dataclass


@dataclass(unsafe_hash=True)
class CartesianPoint:
    x: float
    y: float
    z: float


# creating instance
point = CartesianPoint(0, 0, 0)

# use the class instances as dictionary keys
print({f"{point}": "origin"})

>>> {'CartesianPoint(x=0, y=0, z=0)': 'origin'}

Converting instances to dicts

The asdict() function converts a dataclass instance to a dict of its fields.

from dataclasses import dataclass, asdict

point = CartesianPoint(1, 5, 6)
print(asdict(point))

>>> {'x': 1, 'y': 5, 'z': 6}

Post-init processing

When dataclass generates the __init__ method, internally it’ll call _post_init__ method. You can add additional processing in the __post_init__ method. Here, I’ve added another attribute tup that returns the cartesian point as a tuple.

from dataclasses import dataclass


@dataclass
class CartesianPoint:
    x: float
    y: float
    z: float

    def __post_init__(self):
        self.tup = (self.x, self.y, self.z)


# checking the tuple
point = CartesianPoint(4, 5, 6)
print(point.tup)

>>> (4, 5, 6)

Refactoring the CartesianPoint class

The feature rich original CartesianPoint looks something like this:

import math


class CartesianPoint:
    """Immutable Cartesian point class.
    Although mathematically incorrect,
    for demonstration purpose, all the
    comparisons are done based on
    the first field only."""

    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z

    def __repr__(self):
        """Print the instance neatly."""

        return f"CartesianPoint(x = {self.x}, y = {self.y}, z = {self.z})"

    def __eq__(self, other):
        "Checks if equal."

        return self.x == other.x

    def __nq__(self, other):
        """Checks non equality."""

        return self.x != other.x

    def __gt__(self, other):
        """Checks if greater than."""

        return self.x > other.x

    def __ge__(self, other):
        """Checks if greater than or equal."""

        return self.x >= other.x

    def __lt__(self, other):
        """Checks if less than."""

        return self.x < other.x

    def __le__(self, other):
        """Checks if less than or equal."""

        return self.x <= other.x

    def __hash__(self):
        """Make the instances hashable."""
        return hash(self)

    def dist(self):
        """Finds distance of point from origin."""

        return math.sqrt(self.x**2 + self.y**2 + self.z**2)

Let’s see the class in action:

# create multiple instances of the class
a = CartesianPoint(1, 2, 3)
b = CartesianPoint(1, 3, 3)
c = CartesianPoint(0, 3, 5)
d = CartesianPoint(5, 6, 7)

# checking the __repr__ method
print(a)

# checking the __eq__ method
print(a == b)

# checking the __nq__ method
print(a != c)

# checking the __ge__ method
print(b >= d)

# checking the __lt__ method
print(c < a)

# checking __hash__ and __dist__ method
print({f"{a}": a.dist()})

CartesianPoint(x = 1, y = 2, z = 3)
True
True
False
True
{'CartesianPoint(x = 1, y = 2, z = 3)': 3.7416573867739413}

Below is the same class refactored using dataclass.

from dataclasses import dataclass, field


@dataclass(unsafe_hash=True, order=True)
class CartesianPoint:
    """Immutable Cartesian point class.
    Although mathematically incorrect,
    for demonstration purpose, all the
    comparisons are done based on
    the first field only."""

    x: float
    y: float = field(compare=False)
    z: float = field(compare=False)

    def dist(self):
        """Finds distance of point from origin."""

        return math.sqrt(self.x**2 + self.y**2 + self.z**2)

Use this class like before.

# create multiple instances of the class
a = CartesianPoint(1, 2, 3)
b = CartesianPoint(1, 3, 3)
c = CartesianPoint(0, 3, 5)
d = CartesianPoint(5, 6, 7)

# checking the __repr__ method
print(a)

# checking the __eq__ method
print(a == b)

# checking the __nq__ method
print(a != c)

# checking the __ge__ method
print(b >= d)

# checking the __lt__ method
print(c < a)

# checking __hash__ and __dist__ method
print({f"{a}": a.dist()})

CartesianPoint(x=1, y=2, z=3)
True
True
False
True
{'CartesianPoint(x=1, y=2, z=3)': 3.7416573867739413}

— ⁂ —

Dataclasses#

Using default values#

Using arbitrary field type#

Instance ordering#

Adding methods#

Making instances immutable#

Making instances hashable#

Converting instances to dicts#

Post-init processing#

Refactoring the CartesianPoint class#

Recent posts