Dataclasses, Enums and match 🎁

The boilerplate problem

class PirateManual:
    def __init__(self, name, role, insults=0):
        self.name = name
        self.role = role
        self.insults = insults

    def __repr__(self):
        return f"PirateManual(name={self.name!r}, role={self.role!r}, insults={self.insults!r})"

    def __eq__(self, other):
        if not isinstance(other, PirateManual):
            return NotImplemented
        return (self.name, self.role, self.insults) == (other.name, other.role, other.insults)


print(PirateManual("Guybrush", "captain", 8))

PirateManual(name='Guybrush', role='captain', insults=8)

Fourteen lines, and every field is written three times. Now the same thing:

from dataclasses import dataclass


@dataclass
class Pirate:
    name: str
    role: str
    insults: int = 0


guy = Pirate("Guybrush", "captain", 8)
print(guy)
print(guy.name, guy.insults)
print(Pirate("Elaine", "governor") == Pirate("Elaine", "governor"))

Pirate(name='Guybrush', role='captain', insults=8)
Guybrush 8
True

@dataclass writes __init__, __repr__ and __eq__ from the annotations. The type hints are required (that is how it finds the fields) and, as ever, not enforced at runtime.

Defaults, and the list trap solved properly

from dataclasses import dataclass, field


@dataclass
class Ship:
    name: str
    crew: list[str] = field(default_factory=list)      # a NEW list per instance
    cargo: dict[str, int] = field(default_factory=dict)
    seaworthy: bool = True


a = Ship("Sea Monkey")
b = Ship("Flying Dutchman")

a.crew.append("Otis")
print(a)
print(b)

Ship(name='Sea Monkey', crew=['Otis'], cargo={}, seaworthy=True)
Ship(name='Flying Dutchman', crew=[], cargo={}, seaworthy=True)

Frozen: immutable and hashable

from dataclasses import dataclass


@dataclass(frozen=True)
class Point:
    x: int
    y: int


p = Point(3, 4)
print(p)

try:
    p.x = 10
except Exception as err:
    print(type(err).__name__, err)

# frozen dataclasses are hashable, so they work as keys and in sets
grid = {Point(0, 0): "start", Point(3, 4): "treasure"}
print(grid[Point(3, 4)])
print(len({Point(1, 1), Point(1, 1)}))

Point(x=3, y=4)
FrozenInstanceError cannot assign to field 'x'
treasure
1

frozen=True gives you value semantics: two points with the same coordinates are equal, hash the same, and cannot be modified by a function you passed them to. For coordinates, money, configuration and anything you put in a set, this is what you want.

The other useful options

from dataclasses import dataclass, field, asdict, astuple


@dataclass(order=True, slots=True)
class Score:
    points: int
    name: str = field(compare=False)      # not used when sorting
    notes: str = field(default="", repr=False)


scores = [Score(42, "Otis"), Score(95, "Guybrush"), Score(88, "Elaine")]
print(sorted(scores))
print(max(scores).name)
print(asdict(scores[0]))
print(astuple(scores[0]))

[Score(points=42, name='Otis'), Score(points=88, name='Elaine'), Score(points=95, name='Guybrush')]
Guybrush
{'points': 42, 'name': 'Otis', 'notes': ''}
(42, 'Otis', '')

__post_init__ for validation

from dataclasses import dataclass


@dataclass
class Booking:
    name: str
    seats: int
    price_each: float
    total: float = 0.0

    def __post_init__(self):
        if self.seats < 1:
            raise ValueError(f"seats must be at least 1, got {self.seats}")
        self.total = round(self.seats * self.price_each, 2)


print(Booking("Elaine", 3, 24.99))

try:
    Booking("Otis", 0, 10.0)
except ValueError as err:
    print("Refused:", err)

Booking(name='Elaine', seats=3, price_each=24.99, total=74.97)
Refused: seats must be at least 1, got 0

Enums: named constants that cannot be mistyped

from enum import Enum, auto


class Status(Enum):
    DRAFT = auto()
    PUBLISHED = auto()
    ARCHIVED = auto()


class Suit(Enum):
    HEARTS = "♥"
    SPADES = "♠"


print(Status.DRAFT)
print(Status.DRAFT.name, Status.DRAFT.value)
print(Status("2") if False else Status(2))
print(Suit.HEARTS.value)
print(list(Status))
print(Status.DRAFT == Status.DRAFT, Status.DRAFT == Status.PUBLISHED)

Status.DRAFT
DRAFT 1
Status.PUBLISHED
♥
[<Status.DRAFT: 1>, <Status.PUBLISHED: 2>, <Status.ARCHIVED: 3>]
True False

from enum import StrEnum


class Level(StrEnum):          # Python 3.11+
    DEBUG = "debug"
    INFO = "info"
    ERROR = "error"


print(Level.INFO)
print(Level.INFO == "info")               # behaves as a string too
print(f"level={Level.ERROR}")
print(sorted(Level, key=lambda l: l.value))

info
True
level=error
[<Level.DEBUG: 'debug'>, <Level.ERROR: 'error'>, <Level.INFO: 'info'>]

Notice print(Level.INFO) gave info, not Level.INFO. That is the whole point of StrEnum: it is a string, so it drops straight into f-strings, JSON and database columns while still being a real enum in your code. A plain Enum would have printed Level.INFO and confused whatever you handed it to.

match, properly: structural pattern matching

from dataclasses import dataclass


@dataclass
class Click:
    x: int
    y: int


@dataclass
class KeyPress:
    key: str


@dataclass
class Quit:
    pass


def handle(event):
    match event:
        case Quit():
            return "Goodbye."
        case Click(x=0, y=0):
            return "Clicked the very corner."
        case Click(x=x, y=y) if x == y:
            return f"Clicked the diagonal at {x}."
        case Click(x=x, y=y):
            return f"Clicked at ({x}, {y})."
        case KeyPress(key="q"):
            return "Quit key."
        case KeyPress(key=key):
            return f"Pressed {key!r}."
        case _:
            return "No idea what that was."


for event in [Quit(), Click(0, 0), Click(5, 5), Click(2, 9), KeyPress("q"), KeyPress("a"), 42]:
    print(handle(event))

Goodbye.
Clicked the very corner.
Clicked the diagonal at 5.
Clicked at (2, 9).
Quit key.
Pressed 'a'.
No idea what that was.

This is far more than a switch statement. It matches on shape: the type, the field values, and a guard condition, pulling out the parts you name as it goes. Compare the pile of isinstance checks and attribute lookups you would otherwise write.

Matching data shapes

def describe(data):
    match data:
        case []:
            return "an empty list"
        case [single]:
            return f"one item: {single}"
        case [first, *rest]:
            return f"{first}, then {len(rest)} more"
        case {"type": "user", "name": str(name)}:
            return f"a user called {name}"
        case {"type": kind}:
            return f"some {kind}"
        case str() | bytes():
            return "text of some kind"
        case _:
            return "something else"


for item in [[], [1], [1, 2, 3], {"type": "user", "name": "Elaine"},
             {"type": "ship"}, "hello", 3.14]:
    print(describe(item))

an empty list
one item: 1
1, then 2 more
a user called Elaine
some ship
text of some kind
something else

class Product:
    def __init__(self, name, price, tags=None):
        self.name = name
        self.price = price
        self.tags = tags if tags is not None else []

    def __repr__(self):
        return f"Product({self.name!r}, {self.price!r}, {self.tags!r})"

from dataclasses import dataclass, field


@dataclass(order=True)
class Product:
    price: float
    name: str = field(compare=False)
    tags: list[str] = field(default_factory=list, compare=False)


items = [
    Product(24.99, "Rubber chicken", ["novelty"]),
    Product(4.50, "Grog"),
    Product(12.00, "Map"),
]

for product in sorted(items):
    print(f"{product.price:6.2f}  {product.name}")

  4.50  Grog
 12.00  Map
 24.99  Rubber chicken

def advance(status):
    if status == "draft":
        return "review"
    elif status == "review":
        return "published"
    return status

from enum import Enum


class Status(Enum):
    DRAFT = "draft"
    REVIEW = "review"
    PUBLISHED = "published"


NEXT = {Status.DRAFT: Status.REVIEW, Status.REVIEW: Status.PUBLISHED}


def advance(status: Status) -> Status:
    """Move to the next status, or stay put if already final."""
    return NEXT.get(status, status)


print(advance(Status.DRAFT))
print(advance(Status.REVIEW))
print(advance(Status.PUBLISHED))

try:
    Status("deleted")
except ValueError as err:
    print("Refused:", err)

Status.REVIEW
Status.PUBLISHED
Status.PUBLISHED
Refused: 'deleted' is not a valid Status

def parse(line):
    match line.lower().split():
        case ["look"] | ["l"]:
            return "You see trees. Many trees."
        case ["go", direction]:
            return f"You walk {direction}."
        case ["take", item]:
            return f"Taken: {item}."
        case ["take", count, item] if count.isdigit():
            return f"Taken {count} x {item}."
        case ["say", *words]:
            return f"You say: {' '.join(words)}"
        case []:
            return "Say something."
        case [verb, *_]:
            return f"I do not know how to {verb}."


for line in ["look", "go north", "take grog", "take 3 mints",
             "say you fight like a cow", "", "dance wildly"]:
    print(parse(line))

You see trees. Many trees.
You walk north.
Taken: grog.
Taken 3 x mints.
You say: you fight like a cow
Say something.
I do not know how to dance.