Data Classes
Contents
3.7. Data Classes¶
3.7.1. Data Classes vs Normal Classes¶
If you want to use classes to store data, use the dataclass module. This module is available in Python 3.7+.
With dataclass, you can create a class with attributes, type hints, and a nice representation of the data in a few lines of code. To use dataclass, simply add the @dataclass
decorator on top of a class.
from dataclasses import dataclass
@dataclass
class DataClassDog:
color: str
age: int
DataClassDog(color="black", age=9)
DataClassDog(color='black', age=9)
Without dataclass, you need to use __init__
to assign values to appropriate variables and use __repr__
to create a nice presentation of the data, which can be very cumbersome.
class Dog:
def __init__(self, color, age):
self.color = color
self.age = age
def __repr__(self):
return f"Dog(color={self.color} age={self.age})"
Dog(color="black", age=9)
Dog(color=black age=9)
3.7.2. frozen=True: Make Your Data Classes Read-Only¶
If you don’t want anybody to adjust the attributes of a class, use @dataclass(frozen=True)
.
from dataclasses import dataclass
@dataclass(frozen=True)
class DataClassDog:
color: str
age: int
Now changing the attribute color
of the DataClassDog
’s instance will throw an error.
pepper = DataClassDog(color="black", age=9)
pepper.color = 'golden'
---------------------------------------------------------------------------
FrozenInstanceError Traceback (most recent call last)
<ipython-input-2-0d6f339835b8> in <module>
1 pepper = DataClassDog(color="black", age=9)
----> 2 pepper.color = 'golden'
<string> in __setattr__(self, name, value)
FrozenInstanceError: cannot assign to field 'color'
3.7.3. Compare Between Two Data Classes¶
Normally, you need to implement the __eq__
method so that you can compare between two classes.
class Dog:
def __init__(self, type, age):
self.type = type
self.age = age
def __eq__(self, other):
return (self.type == other.type
and self.age == other.age)
pepper = Dog(type="Dachshund", age=7)
bim = Dog(type="Dachshund", age=7)
pepper == bim
True
dataclasses automatically implements the __eq__
method for you. With dataclasses, you can compare between 2 classes by only specifying their attributes.
from dataclasses import dataclass
@dataclass
class DataClassDog:
type: str
age: int
pepper = DataClassDog(type="Dachshund", age=7)
bim = DataClassDog(type="Dachshund", age=7)
pepper == bim
True
3.7.4. Post-init: Add Init Method to a Data Class¶
With a data class, you don’t need an __init__
method to assign values to its attributes. However, sometimes you might want to use an ___init__
method to initialize certain attributes. That is when data class’s __post_init__
comes in handy.
In the code below, I use __post_init__
to initialize the attribute info
using the attributes names
and ages
.
from dataclasses import dataclass
from typing import List
@dataclass
class Dog:
names: str
age: int
@dataclass
class Dogs:
names: List[str]
ages: List[int]
def __post_init__(self):
self.info = [Dog(name, age) for name, age in zip(self.names, self.ages)]
names = ['Bim', 'Pepper']
ages = [5, 6]
dogs = Dogs(names, ages)
dogs.info
[Dog(names='Bim', age=5), Dog(names='Pepper', age=6)]