16. The data model: making your objects feel built in

Python's syntax is a set of hooks. len(x) calls x.__len__(), x[i] calls x.__getitem__(i), a + b calls a.__add__(b), with x: calls __enter__/__exit__. These are the dunder (double underscore) methods. Implement them and your objects work with every built-in function and syntax, which is how NumPy arrays support + and @, and how PyTorch tensors support slicing.

16.1 Real-world example: a 2-D vector

import math
from dataclasses import dataclass

@dataclass(frozen=True, slots=True)
class Vec:
    x: float
    y: float

    def __add__(self, other: "Vec") -> "Vec":
        return Vec(self.x + other.x, self.y + other.y)

    def __mul__(self, k: float) -> "Vec":           # v * 3
        return Vec(self.x * k, self.y * k)

    __rmul__ = __mul__                              # 3 * v

    def __matmul__(self, other: "Vec") -> float:    # v @ w: dot product
        return self.x * other.x + self.y * other.y

    def __abs__(self) -> float:                     # abs(v)
        return math.hypot(self.x, self.y)

    def __bool__(self) -> bool:                     # if v:
        return bool(self.x or self.y)

    def __iter__(self):                             # x, y = v
        yield self.x
        yield self.y

v, w = Vec(3, 4), Vec(1, 0)
print(v + w, 2 * v, v @ w, abs(v), bool(Vec(0, 0)))
x, y = v
print(x, y)

16.2 Containers: __len__, __getitem__, __contains__

class Dataset:
    """A tiny map-style dataset, the same shape torch.utils.data.Dataset expects."""

    def __init__(self, rows: list[tuple[str, int]]) -> None:
        self._rows = rows

    def __len__(self) -> int:
        return len(self._rows)

    def __getitem__(self, i: int) -> tuple[str, int]:
        return self._rows[i]             # also makes the object iterable and sliceable via the list

ds = Dataset([("refund please", 1), ("love it", 0), ("broken on arrival", 1)])
print(len(ds), ds[0], ds[-1])
for text, label in ds:                   # iteration falls back to __getitem__(0, 1, 2...)
    print(label, text)

That is the entire contract of a PyTorch map-style dataset (Ch 58): two methods. The DataLoader does not care what class you are, only that len() and indexing work. This is duck typing: "if it walks like a duck and quacks like a duck, it is a duck."

16.3 Equality and hashing

If you define __eq__, define __hash__ consistently: equal objects must have equal hashes, or dicts and sets will silently misbehave. A mutable object should not be hashable at all (its hash would change while it sits in a set). @dataclass(frozen=True) generates both correctly; @dataclass with eq=True and not frozen sets __hash__ = None, making instances unhashable, which is the safe choice.

16.4 __repr__ versus __str__

repr(x) is for developers: unambiguous, ideally valid Python (Vec(x=3, y=4)). str(x) is for end users (print, f-strings). If you implement only one, implement __repr__. Your logs and debugger will thank you.

16.5 Context managers as classes

import threading

class Timeout:
    """Raise TimeoutError in the main thread's next check if the block runs too long."""

    def __init__(self, seconds: float) -> None:
        self.seconds = seconds
        self.expired = threading.Event()

    def __enter__(self) -> "Timeout":
        self._timer = threading.Timer(self.seconds, self.expired.set)
        self._timer.start()
        return self

    def __exit__(self, exc_type, exc, tb) -> bool:
        self._timer.cancel()
        return False                  # False: do not swallow exceptions from the block

with Timeout(0.5) as t:
    total = sum(range(100_000))
    print("expired?", t.expired.is_set())

__exit__ receives the exception (if any). Returning True suppresses it, which is almost never what you want.