Advanced Generic Patterns

Predict what happens with this code:

from typing import Self

class QueryBuilder:
    def __init__(self) -> None:
        self._table: str = ""
        self._conditions: list[str] = []

    def table(self, name: str) -> Self:
        self._table = name
        return self

    def where(self, condition: str) -> Self:
        self._conditions.append(condition)
        return self

class PostgresQueryBuilder(QueryBuilder):
    def returning(self, *columns: str) -> Self:
        # Postgres-specific RETURNING clause
        return self

query = PostgresQueryBuilder().table("users").where("age > 18").returning("id")
# What is the type of `query`?

The type of query is PostgresQueryBuilder, not QueryBuilder. Without Self, calling .table() would return QueryBuilder, and the subsequent .returning() call would fail because QueryBuilder has no returning method. Self (PEP 673) preserves the subclass type through the fluent chain. This is just one of several advanced generic patterns that make framework-level code type-safe.

What You Will Learn

• The Self type and how it solves the fluent interface problem
• TypeVarTuple and Unpack for variadic generics (PEP 646)
• Recursive type definitions
• Generic protocols for reusable structural interfaces
• Generic type aliases for complex type expressions
• Real-world patterns: builder APIs, tensor shapes, event systems

Prerequisites

• TypeVar, Generic, and bound/constrained TypeVars from Lesson 1
• Protocol from Lesson 2
• Understanding of method chaining and builder patterns
• Familiarity with *args typing and tuple types

Part 1 -- The Self Type (PEP 673)

The Problem Self Solves

Before Self, returning "the current class" from methods in a hierarchy was painful:

from typing import TypeVar

T = TypeVar("T", bound="Shape")

class Shape:
    def set_color(self: T, color: str) -> T:
        self._color = color
        return self

class Circle(Shape):
    def set_radius(self: T, radius: float) -> T:
        self._radius = radius
        return self

# This works but is verbose and error-prone
c = Circle().set_color("red").set_radius(5.0)

The self: T pattern works but has problems: every method must declare the TypeVar, subclasses might forget to propagate it, and it is ugly.

Self Makes It Clean

from typing import Self

class Shape:
    def __init__(self) -> None:
        self._color: str = "black"

    def set_color(self, color: str) -> Self:
        self._color = color
        return self

    def copy(self) -> Self:
        import copy
        return copy.copy(self)

class Circle(Shape):
    def __init__(self) -> None:
        super().__init__()
        self._radius: float = 1.0

    def set_radius(self, radius: float) -> Self:
        self._radius = radius
        return self

# Full type safety through the chain:
c: Circle = Circle().set_color("blue").set_radius(10.0)
# c is Circle, not Shape -- Self resolves to the actual class

s: Shape = Shape().set_color("red")
# s is Shape

# copy() also returns the right type:
c2: Circle = c.copy()  # Circle, not Shape

Self in Class Methods

from typing import Self
from dataclasses import dataclass

@dataclass
class Config:
    debug: bool = False
    log_level: str = "INFO"
    max_retries: int = 3

    @classmethod
    def from_env(cls) -> Self:
        """Load config from environment variables."""
        import os
        return cls(
            debug=os.getenv("DEBUG", "false").lower() == "true",
            log_level=os.getenv("LOG_LEVEL", "INFO"),
            max_retries=int(os.getenv("MAX_RETRIES", "3")),
        )

    @classmethod
    def development(cls) -> Self:
        return cls(debug=True, log_level="DEBUG", max_retries=1)

class ProductionConfig(Config):
    max_retries: int = 5

dev = ProductionConfig.development()
# dev is ProductionConfig, not Config

Self in enter and exit

from typing import Self
from types import TracebackType

class ManagedResource:
    def __enter__(self) -> Self:
        print("Acquiring resource")
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        print("Releasing resource")

class DatabaseConnection(ManagedResource):
    def query(self, sql: str) -> list[dict[str, object]]:
        return [{"result": "data"}]

# The 'with' statement preserves the subclass type:
with DatabaseConnection() as conn:
    # conn is DatabaseConnection, not ManagedResource
    results = conn.query("SELECT * FROM users")

tip

Use Self whenever a method returns self, cls(...), or a copy of the instance. It replaces the old TypeVar("T", bound="ClassName") pattern and is far more readable.

Available in Python 3.11+. For earlier versions, import from typing_extensions.

Part 2 -- TypeVarTuple and Unpack (PEP 646)

The Problem: Variadic Generics

Regular TypeVar represents a single type. What if you need to represent an arbitrary number of types? Consider a function that zips multiple iterables:

# How do you type this?
def my_zip(iter1, iter2, iter3, ...):
    ...
# The number of arguments varies, and each can have a different type

TypeVarTuple (PEP 646, Python 3.11) captures a variable number of types.

Basic TypeVarTuple

from typing import TypeVarTuple, Unpack

Ts = TypeVarTuple("Ts")

def first_elements(*args: Unpack[tuple[*Ts]]) -> tuple[*Ts]:
    """Return the first element of each iterable."""
    # This is a simplified signature for illustration
    return args

result = first_elements(1, "hello", 3.14)
# result is tuple[int, str, float]

Generic Classes with TypeVarTuple

from typing import TypeVarTuple, Generic, Unpack

Ts = TypeVarTuple("Ts")

class TypedTuple(Generic[*Ts]):
    """A wrapper around a tuple with named access."""
    def __init__(self, *args: Unpack[tuple[*Ts]]) -> None:
        self._data: tuple[*Ts] = args

    @property
    def values(self) -> tuple[*Ts]:
        return self._data

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

record: TypedTuple[str, int, bool] = TypedTuple("Alice", 30, True)
vals = record.values  # tuple[str, int, bool]

Tensor Shape Typing

One of the primary motivations for TypeVarTuple is typing tensor dimensions in ML frameworks:

from typing import TypeVarTuple, Generic, Unpack, NewType

# Define dimension types
Batch = NewType("Batch", int)
Channels = NewType("Channels", int)
Height = NewType("Height", int)
Width = NewType("Width", int)

Shape = TypeVarTuple("Shape")

class Tensor(Generic[*Shape]):
    """A typed tensor with shape information."""

    def __init__(self, *shape: Unpack[tuple[*Shape]]) -> None:
        self._shape = shape

    @property
    def shape(self) -> tuple[*Shape]:
        return self._shape

    def reshape(self, *new_shape: Unpack[tuple[*Shape]]) -> "Tensor[*Shape]":
        return Tensor(*new_shape)

# Create tensors with known shapes:
image: Tensor[Batch, Channels, Height, Width] = Tensor(
    Batch(32), Channels(3), Height(224), Width(224)
)

# Type checker knows the shape:
grayscale: Tensor[Batch, Height, Width] = Tensor(
    Batch(32), Height(224), Width(224)
)

Combining TypeVar with TypeVarTuple

from typing import TypeVar, TypeVarTuple, Generic, Unpack

T = TypeVar("T")
Ts = TypeVarTuple("Ts")

class Prefix(Generic[T, *Ts]):
    """A container with a head element and tail elements of different types."""

    def __init__(self, head: T, *tail: Unpack[tuple[*Ts]]) -> None:
        self.head: T = head
        self.tail: tuple[*Ts] = tail

p = Prefix("leader", 1, 2.0, True)
# p.head is str
# p.tail is tuple[int, float, bool]

note

TypeVarTuple has significant limitations in current type checkers:

• mypy support is still maturing (some patterns work, others do not)
• pyright has better TypeVarTuple support as of 2024
• Complex operations like slicing or mapping over TypeVarTuples are often not fully supported

Use TypeVarTuple for shape typing and simple variadic patterns. Fall back to *args: Any for complex cases where type checker support is incomplete.

Part 3 -- Recursive Types

The Challenge of Recursive Data

JSON-like structures are inherently recursive: a value can be a dict containing more values:

from typing import Union

# This is a recursive type alias:
JsonValue = Union[str, int, float, bool, None, list["JsonValue"], dict[str, "JsonValue"]]

def flatten_json(data: JsonValue, prefix: str = "") -> dict[str, str | int | float | bool | None]:
    """Flatten a nested JSON structure into dot-separated keys."""
    result: dict[str, str | int | float | bool | None] = {}

    if isinstance(data, dict):
        for key, value in data.items():
            new_prefix = f"{prefix}.{key}" if prefix else key
            result.update(flatten_json(value, new_prefix))
    elif isinstance(data, list):
        for i, value in enumerate(data):
            result.update(flatten_json(value, f"{prefix}[{i}]"))
    else:
        result[prefix] = data

    return result

# Type-safe usage:
config: JsonValue = {
    "database": {
        "host": "localhost",
        "port": 5432,
        "replicas": [
            {"host": "replica1", "port": 5433},
            {"host": "replica2", "port": 5434},
        ],
    },
    "debug": True,
}

flat = flatten_json(config)
# {"database.host": "localhost", "database.port": 5432, ...}

Recursive Generic Types: Tree Structures

from __future__ import annotations
from typing import Generic, TypeVar, Iterator
from dataclasses import dataclass

T = TypeVar("T")

@dataclass
class TreeNode(Generic[T]):
    value: T
    children: list[TreeNode[T]]

    def __iter__(self) -> Iterator[T]:
        """Depth-first traversal."""
        yield self.value
        for child in self.children:
            yield from child

    def map(self, func: "Callable[[T], U]") -> "TreeNode[U]":
        """Apply a function to every node."""
        return TreeNode(
            value=func(self.value),
            children=[child.map(func) for child in self.children],
        )

    @property
    def depth(self) -> int:
        if not self.children:
            return 1
        return 1 + max(child.depth for child in self.children)

from typing import Callable

U = TypeVar("U")

# Build a tree:
org = TreeNode("CEO", [
    TreeNode("CTO", [
        TreeNode("Lead Engineer", []),
        TreeNode("Senior Engineer", []),
    ]),
    TreeNode("CFO", [
        TreeNode("Accountant", []),
    ]),
])

# Type-safe iteration:
for title in org:
    print(title)  # title is str

# Type-safe mapping:
lengths: TreeNode[int] = org.map(len)
# lengths.value is int

Recursive Protocol

from typing import Protocol, Iterator, Self

class TreeLike(Protocol):
    """Protocol for any tree-like structure."""
    @property
    def value(self) -> object: ...

    @property
    def children(self) -> Iterator[Self]: ...

def count_nodes(tree: TreeLike) -> int:
    return 1 + sum(count_nodes(child) for child in tree.children)

Part 4 -- Generic Type Aliases

Simple Generic Aliases

Type aliases reduce repetition and improve readability:

from typing import TypeVar

T = TypeVar("T")

# Generic type alias:
Result = T | Exception
# Usage: Result[int] means int | Exception

Matrix = list[list[T]]
# Usage: Matrix[float] means list[list[float]]

Callback = "Callable[[T], None]"
# Usage: Callback[str] means Callable[[str], None]

TypeAlias for Explicit Aliases

from typing import TypeAlias, TypeVar

T = TypeVar("T")

# Explicit type alias (Python 3.10+)
JsonPrimitive: TypeAlias = str | int | float | bool | None
JsonArray: TypeAlias = list["JsonValue"]
JsonObject: TypeAlias = dict[str, "JsonValue"]
JsonValue: TypeAlias = JsonPrimitive | JsonArray | JsonObject

# Parameterized alias using TypeVar
Response: TypeAlias = tuple[int, T, dict[str, str]]

def make_response(status: int, body: T, headers: dict[str, str]) -> Response[T]:
    return (status, body, headers)

r = make_response(200, {"users": ["alice"]}, {"content-type": "application/json"})
# r is tuple[int, dict[str, list[str]], dict[str, str]]

The type Statement (Python 3.12+)

Python 3.12 introduces a cleaner syntax for type aliases:

# Python 3.12+ syntax:
type Vector[T] = list[T]
type Matrix[T] = list[Vector[T]]
type Result[T] = T | Exception
type Handler[**P, R] = Callable[P, R]

# Equivalent to:
# T = TypeVar("T")
# Vector = list[T]  # as TypeAlias

note

The type statement is Python 3.12+. For compatibility with earlier versions, use TypeAlias from typing (3.10+) or typing_extensions.

Complex Generic Aliases in Practice

from typing import TypeVar, TypeAlias, Callable, Awaitable

T = TypeVar("T")
E = TypeVar("E", bound=Exception)

# A result type that distinguishes success from failure:
Ok: TypeAlias = tuple[True, T]
Err: TypeAlias = tuple[False, E]
Result: TypeAlias = Ok[T] | Err[E]

def divide(a: float, b: float) -> Result[float, ZeroDivisionError]:
    if b == 0:
        return (False, ZeroDivisionError("division by zero"))
    return (True, a / b)

match divide(10, 3):
    case (True, value):
        print(f"Result: {value}")  # value is float
    case (False, error):
        print(f"Error: {error}")   # error is ZeroDivisionError

# Async handler alias:
AsyncHandler: TypeAlias = Callable[[T], Awaitable[Result[T, Exception]]]

Part 5 -- Multiple TypeVars and Advanced Constraints

Multiple TypeVars in a Single Class

from typing import TypeVar, Generic

K = TypeVar("K")
V = TypeVar("V")

class BiMap(Generic[K, V]):
    """Bidirectional map: look up by key or value."""

    def __init__(self) -> None:
        self._forward: dict[K, V] = {}
        self._backward: dict[V, K] = {}

    def put(self, key: K, value: V) -> None:
        self._forward[key] = value
        self._backward[value] = key

    def get_by_key(self, key: K) -> V | None:
        return self._forward.get(key)

    def get_by_value(self, value: V) -> K | None:
        return self._backward.get(value)

    def items(self) -> list[tuple[K, V]]:
        return list(self._forward.items())

codes = BiMap[str, int]()
codes.put("OK", 200)
codes.put("NOT_FOUND", 404)

print(codes.get_by_key("OK"))        # 200 (type: int | None)
print(codes.get_by_value(404))       # "NOT_FOUND" (type: str | None)

Bound TypeVars with Multiple Constraints

from typing import TypeVar, Protocol

class Measurable(Protocol):
    def measure(self) -> float: ...

class Labelable(Protocol):
    @property
    def label(self) -> str: ...

class MeasurableAndLabelable(Measurable, Labelable, Protocol):
    ...

M = TypeVar("M", bound=MeasurableAndLabelable)

def report(items: list[M]) -> dict[str, float]:
    return {item.label: item.measure() for item in items}

class Sensor:
    def __init__(self, name: str, reading: float) -> None:
        self._name = name
        self._reading = reading

    @property
    def label(self) -> str:
        return self._name

    def measure(self) -> float:
        return self._reading

sensors = [Sensor("temp", 22.5), Sensor("humidity", 65.0)]
data = report(sensors)  # {"temp": 22.5, "humidity": 65.0}

Part 6 -- Real-World Patterns

Pattern: Type-Safe Builder

from typing import Self, Generic, TypeVar
from dataclasses import dataclass, field

@dataclass
class HttpRequest:
    method: str = "GET"
    url: str = ""
    headers: dict[str, str] = field(default_factory=dict)
    body: bytes = b""
    timeout: float = 30.0

class RequestBuilder:
    def __init__(self) -> None:
        self._method: str = "GET"
        self._url: str = ""
        self._headers: dict[str, str] = {}
        self._body: bytes = b""
        self._timeout: float = 30.0

    def method(self, m: str) -> Self:
        self._method = m
        return self

    def url(self, u: str) -> Self:
        self._url = u
        return self

    def header(self, key: str, value: str) -> Self:
        self._headers[key] = value
        return self

    def body(self, b: bytes) -> Self:
        self._body = b
        return self

    def timeout(self, t: float) -> Self:
        self._timeout = t
        return self

    def build(self) -> HttpRequest:
        if not self._url:
            raise ValueError("URL is required")
        return HttpRequest(
            method=self._method,
            url=self._url,
            headers=self._headers,
            body=self._body,
            timeout=self._timeout,
        )

class AuthenticatedRequestBuilder(RequestBuilder):
    def __init__(self) -> None:
        super().__init__()
        self._token: str = ""

    def bearer_token(self, token: str) -> Self:
        self._token = token
        self.header("Authorization", f"Bearer {token}")
        return self

# Full fluent chain with correct types:
req = (
    AuthenticatedRequestBuilder()
    .method("POST")                    # Returns AuthenticatedRequestBuilder
    .url("https://api.example.com")    # Returns AuthenticatedRequestBuilder
    .bearer_token("secret")           # Returns AuthenticatedRequestBuilder
    .header("Content-Type", "json")    # Returns AuthenticatedRequestBuilder
    .body(b'{"key": "value"}')        # Returns AuthenticatedRequestBuilder
    .build()                           # Returns HttpRequest
)

Pattern: Generic Event System

from typing import TypeVar, Generic, Callable
from dataclasses import dataclass

T = TypeVar("T")

class EventBus:
    """Type-safe event bus using generic event types."""

    def __init__(self) -> None:
        self._handlers: dict[type, list[Callable]] = {}

    def subscribe(self, event_type: type[T], handler: Callable[[T], None]) -> None:
        self._handlers.setdefault(event_type, []).append(handler)

    def publish(self, event: T) -> None:
        event_type = type(event)
        for handler in self._handlers.get(event_type, []):
            handler(event)

@dataclass
class UserCreated:
    user_id: str
    email: str

@dataclass
class OrderPlaced:
    order_id: str
    total: float

bus = EventBus()

# Type-safe subscriptions:
def on_user_created(event: UserCreated) -> None:
    print(f"Welcome {event.email}")

def on_order_placed(event: OrderPlaced) -> None:
    print(f"Order {event.order_id}: ${event.total}")

bus.subscribe(UserCreated, on_user_created)  # OK
bus.subscribe(OrderPlaced, on_order_placed)  # OK
# bus.subscribe(UserCreated, on_order_placed)  # Type error if checked

bus.publish(UserCreated("u-1", "[email protected]"))
bus.publish(OrderPlaced("o-1", 99.99))

Pattern: Generic Middleware Chain

from typing import TypeVar, Generic, Callable, Self
from dataclasses import dataclass

In = TypeVar("In")
Out = TypeVar("Out")
Mid = TypeVar("Mid")

class Pipe(Generic[In, Out]):
    """A composable transformation pipeline."""

    def __init__(self, fn: Callable[[In], Out]) -> None:
        self._fn = fn

    def __call__(self, value: In) -> Out:
        return self._fn(value)

    def then(self, fn: Callable[[Out], Mid]) -> "Pipe[In, Mid]":
        prev = self._fn
        return Pipe(lambda x: fn(prev(x)))

    @staticmethod
    def identity() -> "Pipe[In, In]":
        return Pipe(lambda x: x)

# Build a pipeline:
pipeline: Pipe[str, bool] = (
    Pipe(str.strip)                # str -> str
    .then(str.lower)               # str -> str
    .then(lambda s: len(s))        # str -> int
    .then(lambda n: n > 5)         # int -> bool
)

print(pipeline("  Hello World  "))  # True (length 11 > 5)
print(pipeline("  Hi  "))           # False (length 2 > 5 is False)

Part 7 -- Practical Considerations

When to Reach for Advanced Generics

Avoid Over-Generification

danger

Not every function needs generics. If a function works with one or two specific types, use those types directly. Generics add complexity and should be justified by genuine reuse.

# OVER-ENGINEERED:
T = TypeVar("T", bound=str)
def greet(name: T) -> str:
    return f"Hello, {name}"

# JUST RIGHT:
def greet(name: str) -> str:
    return f"Hello, {name}"

Reserve advanced generic patterns for:

• Library code used by many consumers
• Framework abstractions (ORMs, event systems, pipelines)
• Type-safe builder/fluent APIs
• Data structures that genuinely parameterize over element types

Version Compatibility Reference

Feature	Stdlib typing	typing_extensions
Self	3.11+	3.8+
TypeVarTuple	3.11+	3.8+
Unpack	3.11+	3.8+
type statement	3.12+	N/A (syntax)
TypeAlias	3.10+	3.8+

Key Takeaways

• Self (PEP 673) replaces the TypeVar("T", bound="ClassName") pattern for methods returning self or cls(...) -- use it for fluent APIs, builders, context managers, and classmethod constructors
• TypeVarTuple (PEP 646) captures a variable number of types -- use it for tensor shapes, generic tuples, and variadic function types
• Recursive types are expressed through self-referencing type aliases -- essential for trees, JSON structures, and nested data
• Generic type aliases reduce repetition for complex type expressions and can be parameterized with TypeVars
• Generic protocols combine structural subtyping with parametric polymorphism for maximum flexibility
• Use the type statement (Python 3.12+) for cleaner alias syntax
• Apply advanced generics judiciously -- premature generification is as harmful as premature optimization
• Always check type checker support for newer features, especially TypeVarTuple

Graded Practice Challenges

Level 1 -- Predict the Type Checker Output

Question 1: What is the type of result?

from typing import Self

class Node:
    def __init__(self, value: int) -> None:
        self.value = value

    def double(self) -> Self:
        self.value *= 2
        return self

class SpecialNode(Node):
    def triple(self) -> Self:
        self.value *= 3
        return self

result = SpecialNode(5).double().triple()

Answer

result is SpecialNode. Self in double() resolves to the actual class of the instance. Since we start with SpecialNode(5), double() returns SpecialNode, and triple() also returns SpecialNode. Without Self, double() would return Node, and the .triple() call would fail.

Question 2: Does this recursive type alias work with mypy?

from typing import TypeAlias

Nested: TypeAlias = int | list["Nested"]

def flatten(data: Nested) -> list[int]:
    if isinstance(data, int):
        return [data]
    result: list[int] = []
    for item in data:
        result.extend(flatten(item))
    return result

flatten([1, [2, [3, 4]], 5])

Answer

Yes, this works with mypy. Recursive type aliases using forward references (strings) are supported. mypy correctly narrows data to int in the isinstance branch and to list[Nested] in the else branch, where iteration and recursive calls are valid.

Question 3: What is wrong with this code?

from typing import Self

class Singleton:
    _instance: "Singleton | None" = None

    def __new__(cls) -> Self:
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance  # Is this valid?

Answer

This is technically valid but has a subtle issue. cls._instance is typed as Singleton | None, and after the None check, it is narrowed to Singleton. However, Self means "the type of cls", which could be a subclass. If SubSingleton(Singleton) is created, cls._instance could hold a Singleton instance being returned as SubSingleton. mypy may flag this depending on strictness. The fix is to type _instance as Self | None, but class variables cannot use Self directly. This is a known limitation of the Singleton pattern with type-safe inheritance.

Level 2 -- Debug and Fix

This generic linked list has type errors. Find and fix them:

from typing import Generic, TypeVar, Iterator, Self

T = TypeVar("T")

class LinkedList(Generic[T]):
    def __init__(self) -> None:
        self._head: _Node[T] | None = None

    def prepend(self, value: T) -> Self:
        self._head = _Node(value, self._head)
        return self

    def __iter__(self) -> Iterator[T]:
        current = self._head
        while current:
            yield current.value
            current = current.next

    def map(self, func: Callable[[T], U]) -> LinkedList[U]:
        result = LinkedList()
        for value in self:
            result.prepend(func(value))
        return result

class _Node(Generic[T]):
    def __init__(self, value: T, next: "_Node[T] | None" = None) -> None:
        self.value = value
        self.next = next

from typing import Callable
U = TypeVar("U")

Answer

Several issues:

1. Missing imports at the top: Callable and U are used before they are imported/defined. Move them to the top.

2. map returns LinkedList[U] but uses Self-returning prepend: When map creates a new LinkedList(), it is LinkedList[U]. But prepend returns Self, and since map calls it on a plain LinkedList, this works. However, the items will be in reverse order.

3. _Node is referenced before definition: In LinkedList.__init__, _Node[T] is used but _Node is defined later. Use from __future__ import annotations or string annotations.

Fix:

from __future__ import annotations
from typing import Generic, TypeVar, Iterator, Self, Callable

T = TypeVar("T")
U = TypeVar("U")

class _Node(Generic[T]):
    def __init__(self, value: T, next: _Node[T] | None = None) -> None:
        self.value = value
        self.next = next

class LinkedList(Generic[T]):
    def __init__(self) -> None:
        self._head: _Node[T] | None = None

    def prepend(self, value: T) -> Self:
        self._head = _Node(value, self._head)
        return self

    def __iter__(self) -> Iterator[T]:
        current = self._head
        while current:
            yield current.value
            current = current.next

    def map(self, func: Callable[[T], U]) -> LinkedList[U]:
        result: LinkedList[U] = LinkedList()
        items = list(self)
        for value in reversed(items):  # Reverse to maintain order
            result.prepend(func(value))
        return result

Level 3 -- Design Challenge

Design a type-safe state machine where:

# States and transitions are checked at type level:
class Draft: pass
class Review: pass
class Published: pass

machine = StateMachine[Draft]()  # Start in Draft state

# Only valid transitions should type-check:
machine.transition(Review)      # OK: Draft -> Review
machine.transition(Published)   # OK: Review -> Published
machine.transition(Draft)       # Should ERROR: Published -> Draft not allowed

Requirements:

1. Use generics to track the current state type
2. Transitions return a new StateMachine with the target state type
3. Invalid transitions should be caught by the type checker
4. Support attaching callbacks to transitions

Hint

One approach uses method overloads per valid transition:

from typing import Generic, TypeVar, Callable, overload

S = TypeVar("S")

class StateMachine(Generic[S]):
    def __init__(self, state: S) -> None:
        self._state = state
        self._callbacks: list[Callable] = []

    @property
    def state(self) -> S:
        return self._state

class Draft:
    pass

class Review:
    pass

class Published:
    pass

class DocumentMachine(StateMachine[Draft]):
    @overload
    def transition(self: "DocumentMachine", to: type[Review]) -> "ReviewMachine": ...
    @overload
    def transition(self: "DocumentMachine", to: type[Published]) -> None: ...

    def transition(self, to):
        if to is Review:
            return ReviewMachine(Review())
        raise ValueError(f"Cannot transition from Draft to {to}")

class ReviewMachine(StateMachine[Review]):
    def transition(self, to: type[Published]) -> StateMachine[Published]:
        return StateMachine(Published())

A more scalable approach defines allowed transitions as a mapping and validates them. See the full design challenge for exploration.

What's Next

In the next lesson, Overload and Type Narrowing, we explore how @overload gives functions multiple type signatures and how TypeGuard, TypeIs, and assert_never enable exhaustive type narrowing.