Dataclasses & Enums

Two standard-library tools remove a lot of boilerplate. Dataclasses generate __init__, __repr__ and __eq__ for classes that mainly hold data. Enums define a fixed set of named constants, so a value like a status or a colour can only be one of the allowed options.

A basic dataclass

PY dataclass.py
from dataclasses import dataclass

@dataclass
class Book:
    title: str
    author: str
    pages: int
    available: bool = True

b = Book("Dune", "Frank Herbert", 412)
print(b)
print(b == Book("Dune", "Frank Herbert", 412))
b.available = False
print(b.available)
Output
Book(title='Dune', author='Frank Herbert', pages=412, available=True)
True
False

From the annotated class attributes, @dataclass writes the initialiser, a readable repr and value-based equality for you. Fields with defaults must come after fields without them.

Mutable defaults and computed fields

PY fields.py
from dataclasses import dataclass, field

@dataclass
class Order:
    customer: str
    items: list[str] = field(default_factory=list)
    total: float = 0.0
    order_id: int = field(default=0, repr=False)

    def __post_init__(self):
        self.customer = self.customer.strip().title()

a = Order("  ana silva ")
b = Order("ben")
a.items.append("tea")
print(a)
print(b.items)
Output
Order(customer='Ana Silva', items=['tea'], total=0.0)
[]

Lists and dicts can't be plain defaults (the mutable-default trap from lesson 29); field(default_factory=list) creates a fresh one per object. __post_init__ runs after the generated __init__, which is handy for cleaning or validating data.

Frozen, ordered and slotted dataclasses

PY options.py
from dataclasses import dataclass, replace

@dataclass(frozen=True, order=True, slots=True)
class Version:
    major: int
    minor: int
    patch: int = 0

v = Version(1, 4)
print(sorted([Version(2, 0), v, Version(1, 10)]))
print(replace(v, patch=1))
print(len({Version(1, 0), Version(1, 0)}))

try:
    v.major = 9
except Exception as err:
    print(type(err).__name__)
Output
[Version(major=1, minor=4, patch=0), Version(major=1, minor=10, patch=0), Version(major=2, minor=0, patch=0)]
Version(major=1, minor=4, patch=1)
1
FrozenInstanceError
  • frozen=True makes instances immutable (and hashable, so usable in sets and as dict keys). Use replace() to get a modified copy.
  • order=True generates comparison methods, comparing fields in order.
  • slots=True uses less memory and prevents adding unexpected attributes.

Converting to dictionaries

PY asdict.py
from dataclasses import dataclass, asdict
import json

@dataclass
class Point:
    x: float
    y: float

p = Point(1.5, -2)
print(asdict(p))
print(json.dumps(asdict(p)))
Output
{'x': 1.5, 'y': -2}
{"x": 1.5, "y": -2}

Enums

PY enum_basic.py
from enum import Enum, auto

class Status(Enum):
    PENDING = auto()
    SHIPPED = auto()
    DELIVERED = auto()

s = Status.SHIPPED
print(s, s.name, s.value)
print(s is Status.SHIPPED)
print([member.name for member in Status])
print(Status["DELIVERED"], Status(1))
Output
Status.SHIPPED SHIPPED 2
True
['PENDING', 'SHIPPED', 'DELIVERED']
Status.DELIVERED Status.PENDING

An enum's members are unique, named constants. You can look them up by name (Status["DELIVERED"]) or by value (Status(1)), and loop over all of them. auto() assigns values automatically when the actual value doesn't matter.

Why enums beat strings

PY enum_vs_string.py
from enum import StrEnum

class Color(StrEnum):
    RED = "red"
    GREEN = "green"
    BLUE = "blue"

def paint(color: Color):
    return f"Painting it {color}"

print(paint(Color.GREEN))
print(Color("blue") is Color.BLUE)

try:
    Color("purple")
except ValueError as err:
    print(err)
Output
Painting it green
True
'purple' is not a valid Color

With plain strings, a typo like "gren" silently flows through your program. With an enum, invalid values are rejected immediately and editors autocomplete the valid ones. StrEnum members also behave as strings, which is convenient for JSON and databases; IntEnum does the same for numbers.

Enums with behaviour

PY enum_methods.py
from enum import Enum

class Planet(Enum):
    MERCURY = (3.303e23, 2.4397e6)
    EARTH = (5.976e24, 6.37814e6)

    def __init__(self, mass, radius):
        self.mass = mass
        self.radius = radius

    @property
    def surface_gravity(self):
        G = 6.673e-11
        return G * self.mass / self.radius ** 2

for p in Planet:
    print(f"{p.name.title():<8} {p.surface_gravity:.2f} m/s²")
Output
Mercury  3.70 m/s²
Earth    9.80 m/s²

Enums and match

PY enum_match.py
from enum import Enum, auto

class Light(Enum):
    RED = auto()
    AMBER = auto()
    GREEN = auto()

def action(light):
    match light:
        case Light.RED:
            return "stop"
        case Light.AMBER:
            return "prepare"
        case Light.GREEN:
            return "go"

print([action(l) for l in Light])
Output
['stop', 'prepare', 'go']

Try it yourself

Try it yourself
Console output
PreviousProperties & Class Methods