Object-oriented Python · about 45 min

Dataclasses and enums

Why this matters

A lot of classes are mostly data: a point, a config entry, a row from a CSV file, a message. Writing __init__, __repr__ and __eq__ by hand for each is tedious and a source of bugs (forget one field in __eq__ and two different objects compare equal). The dataclasses module generates these methods from a list of typed fields. And when a value can only be one of a fixed set of choices (a traffic light color, an HTTP method, an order status), the enum module gives you named constants that cannot be misspelled. Both are used everywhere in modern Python.

From hand-written to @dataclass

The manual version of a simple data class:

class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y
    def __repr__(self):
        return f"Point(x={self.x!r}, y={self.y!r})"
    def __eq__(self, other):
        if not isinstance(other, Point):
            return NotImplemented
        return (self.x, self.y) == (other.x, other.y)

The dataclass version:

from dataclasses import dataclass

@dataclass
class Point:
    x: float
    y: float

p = Point(1.5, 2)
print(p)
print(p == Point(1.5, 2))
print(Point(y=0, x=3))

Point(x=1.5, y=2)
True
Point(x=3, y=0)

x: float is a variable annotation: a name with a type hint and no value. The @dataclass decorator runs once, when the class is created. It reads the class's annotations (Point.__annotations__, in order), then generates source code for __init__, __repr__ and __eq__ and attaches the resulting methods to the class. There is no magic at attribute-access time; the resulting class is an ordinary class you could have written yourself.

The type hints are not enforced: Point("a", None) works fine. They document intent and are used by tools such as type checkers (the type hints lesson in the next module covers this).

Defaults and the mutable default problem

Fields can have defaults; as with function parameters, fields with defaults must come after those without:

@dataclass
class Task:
    title: str
    done: bool = False
    priority: int = 3

print(Task("write lesson"))

Task(title='write lesson', done=False, priority=3)

Mutable defaults would be shared between instances (the same trap as mutable default arguments and mutable class attributes), so dataclasses refuse them:

@dataclass
class Bag:
    items: list = []

ValueError: mutable default <class 'list'> for field items is not allowed: use default_factory

Use field(default_factory=...): a function called to build a fresh default for every instance:

from dataclasses import dataclass, field

@dataclass
class Bag:
    items: list[str] = field(default_factory=list)

a, b = Bag(), Bag()
a.items.append("key")
print(a, b)

Bag(items=['key']) Bag(items=[])

field() takes other options too, for example field(repr=False) to hide a field (say, a password) from the repr, or field(compare=False) to leave it out of equality.

Validation with post_init

The generated __init__ just assigns fields. If you need checks or derived values, define __post_init__, which the generated __init__ calls at the end:

@dataclass
class Rect:
    w: float
    h: float
    area: float = field(init=False)     # not a parameter; computed

    def __post_init__(self):
        if self.w < 0 or self.h < 0:
            raise ValueError("sizes must be non-negative")
        self.area = self.w * self.h

print(Rect(2, 3))

Rect(w=2, h=3, area=6)

Options: frozen and order

@dataclass(...) takes arguments that change what is generated:

order=True adds __lt__, __le__, __gt__, __ge__, comparing the fields as a tuple in declaration order.
frozen=True makes instances immutable: assigning a field raises FrozenInstanceError. Because a frozen object's fields can never change, the dataclass also generates __hash__ from the fields, so instances can be dict keys and set members.

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

v = Version(1, 2)
print(sorted([Version(2, 0), v, Version(1, 10)]))
print({v: "current"}[Version(1, 2)])
v.major = 3

[Version(major=1, minor=2), Version(major=1, minor=10), Version(major=2, minor=0)]
current
dataclasses.FrozenInstanceError: cannot assign to field 'major'

Why does a plain dataclass have no hash? Recall from lesson 2 that defining __eq__ sets __hash__ to None. A dict finds a key by its hash; if a mutable object's fields (and so its hash) changed after insertion, the dict would look in the wrong place and lose it. Frozen objects cannot change, so hashing them is safe. The internals module explains hash tables fully.

Other useful options: slots=True stores fields in fixed slots instead of a per-instance __dict__ (less memory, faster access, no adding stray attributes), and kw_only=True requires keyword arguments. Helper functions: dataclasses.asdict(obj) converts to a dict (handy for JSON), and dataclasses.replace(obj, field=new) makes a modified copy, the idiomatic way to "change" a frozen instance.

Enums: a fixed set of named values

Using bare strings for states ("red", "gren"...) means typos are silently accepted. An Enum defines the complete set once:

from enum import Enum

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

c = Color.GREEN
print(c, c.name, c.value)
print(Color("blue"))           # look up by value
print(Color["RED"])            # look up by name
print([m.name for m in Color]) # iteration, in definition order
print(c is Color.GREEN)
Color("purple")

Color.GREEN GREEN green
Color.BLUE
Color.RED
['RED', 'GREEN', 'BLUE']
True
ValueError: 'purple' is not a valid Color

Each member is a unique singleton object, an instance of the Color class, so compare with is (or ==). Enum classes are classes, so they can have methods:

class Status(Enum):
    TODO = 1
    DOING = 2
    DONE = 3

    def is_open(self):
        return self is not Status.DONE

print(Status.DOING.is_open(), Status.DONE.is_open())

True False

When you do not care about the values, enum.auto() assigns them (1, 2, 3 ...). IntEnum and StrEnum members also behave as ints or strings, useful when interfacing with code that expects plain values. Enums combine nicely with match statements: case Status.DONE:.

Common mistakes

Mutable default without default_factory: the ValueError shown above.
Non-default field after a default one: TypeError: non-default argument 'y' follows default argument 'x'.
Expecting type checking: annotations are not enforced; validate in __post_init__ if it matters.
Mutating a frozen instance in __post_init__: even there, self.x = ... raises FrozenInstanceError (use object.__setattr__(self, "x", ...) if you really must).
Comparing enum members to raw values: Color.RED == "red" is False for a plain Enum. Compare to Color.RED, or convert with Color("red").

Recap

@dataclass generates __init__, __repr__ and __eq__ from annotated fields; it is code generation at class creation time.
Use field(default_factory=list) for mutable defaults and __post_init__ for validation and derived fields.
frozen=True makes instances immutable and hashable; order=True adds comparisons by field tuple.
Enum gives a closed set of named singleton members with .name and .value, lookup by value Cls(v) or name Cls[n].

Going further: in the REPL, run help(Point) on a dataclass to see the signatures of the generated methods, and try dataclasses.fields(Point). Read the dataclasses and enum docs ("Enum HOWTO" is excellent).

Exercises

Exercise 1 · quiz

What @dataclass writes for you

Which methods does a plain @dataclass (no arguments) generate? Select all that apply.

Answer

init
repr
eq

By default you get __init__, __repr__ and __eq__. Ordering methods need order=True. Because __eq__ is generated, __hash__ is set to None (unhashable) unless you ask for frozen=True (or unsafe_hash=True), in which case a hash based on the fields is generated: equal objects must have equal hashes, and mutable objects that can change their hash would break dicts and sets.

Exercise 2 · write code

An inventory with dataclasses

In solution.py, using from dataclasses import dataclass, field:

A dataclass Item with fields name: str, price: float, quantity: int = 0. Add a __post_init__ that raises ValueError if price is negative or quantity is negative.
A dataclass Inventory with a field items: list[Item] that defaults to a new empty list for each inventory (use field(default_factory=list)), and methods:
- add(item) appends it
- total_value() returns the sum of price * quantity over all items
- names() returns the item names sorted alphabetically

The generated repr(Item("pen", 1.5, 10)) must be Item(name='pen', price=1.5, quantity=10).

File: solution.py

Starting point

from dataclasses import dataclass, field


@dataclass
class Item:
    name: str
    price: float
    quantity: int = 0


@dataclass
class Inventory:
    items: list = []

Hint 1

items: list = [] is rejected by dataclasses (ValueError: mutable default ... use default_factory). Write field(default_factory=list).

Hint 2

__post_init__(self) runs right after the generated __init__; check self.price and self.quantity there.

Hint 3

sum(i.price * i.quantity for i in self.items) computes the total.

Reference solution

from dataclasses import dataclass, field


@dataclass
class Item:
    name: str
    price: float
    quantity: int = 0

    def __post_init__(self):
        if self.price < 0:
            raise ValueError(f"negative price: {self.price}")
        if self.quantity < 0:
            raise ValueError(f"negative quantity: {self.quantity}")


@dataclass
class Inventory:
    items: list[Item] = field(default_factory=list)

    def add(self, item):
        self.items.append(item)

    def total_value(self):
        return sum(i.price * i.quantity for i in self.items)

    def names(self):
        return sorted(i.name for i in self.items)

Exercise 3 · write code

Frozen, ordered coordinates

In solution.py, write a dataclass Coord with integer fields row and col, declared with frozen=True and order=True, so that:

coordinates are immutable (assigning c.row = 5 raises dataclasses.FrozenInstanceError),
they can be dict keys and set members,
they sort by row, then col.

Also write a function neighbors(c) that returns a sorted list of the four Coords directly above, below, left and right of c.

Example: neighbors(Coord(1, 1)) returns [Coord(row=0, col=1), Coord(row=1, col=0), Coord(row=1, col=2), Coord(row=2, col=1)].

File: solution.py

Starting point

from dataclasses import dataclass


@dataclass
class Coord:
    row: int
    col: int


def neighbors(c):
    pass

Hint 1

Pass options to the decorator: @dataclass(frozen=True, order=True).

Hint 2

With order=True, fields are compared as a tuple in declaration order, so sorted just works.

Reference solution

from dataclasses import dataclass


@dataclass(frozen=True, order=True)
class Coord:
    row: int
    col: int


def neighbors(c):
    return sorted([
        Coord(c.row - 1, c.col),
        Coord(c.row + 1, c.col),
        Coord(c.row, c.col - 1),
        Coord(c.row, c.col + 1),
    ])

Exercise 4 · write code

An Enum state machine

In solution.py, write an Enum called Light with members RED, GREEN, YELLOW (values "red", "green", "yellow"), and a method next(self) that returns the following state: RED -> GREEN -> YELLOW -> RED.

Then write a function run(start, steps) where start is a string like "red". It converts the string to a Light with Light(start) (an unknown value then raises ValueError automatically) and returns a list of the member names visited, starting with the start state and taking steps transitions.

Example: run("red", 4) returns ["RED", "GREEN", "YELLOW", "RED", "GREEN"].

File: solution.py

Starting point

from enum import Enum

RED, GREEN, YELLOW = "red", "green", "yellow"


def run(start, steps):
    pass

Hint 1

class Light(Enum): with RED = "red" and so on. Methods go in the class body as usual.

Hint 2

Light("red") looks a member up by value; member.name is "RED".

Hint 3

For next, a list of the members in cycle order and (index + 1) % 3 does the wrap-around.

Reference solution

from enum import Enum


class Light(Enum):
    RED = "red"
    GREEN = "green"
    YELLOW = "yellow"

    def next(self):
        order = [Light.RED, Light.GREEN, Light.YELLOW]
        return order[(order.index(self) + 1) % len(order)]


def run(start, steps):
    state = Light(start)
    names = [state.name]
    for _ in range(steps):
        state = state.next()
        names.append(state.name)
    return names