Type Hints, Generics, Protocols

TL;DR

Python’s runtime is dynamically typed and that won’t change. Type hints are annotations — they’re not enforced at runtime. Their value is in the editor (autocomplete, jump-to-def) and in a separate type checker (mypy or pyright) that reads the annotations and tells you when types don’t line up.

The reason to use them in ML: codebases get big. By the time you have 12 files calling each other and a train() function whose signature you keep changing, “what does this function expect?” stops being obvious. Type hints answer the question without you reading the body.

The minimum useful set:

Use pyright (faster, stricter, what VS Code’s Python extension runs) or mypy (older, more configurable, more in-tree integrations). Either is fine; pick one.

The picture in your head

A type hint is a comment that a tool reads. The Python interpreter treats it as decoration on the AST and ignores it for execution. mypy / pyright reads the same annotations and traces them through your code: “you said train(loader: DataLoader, model: Module) -> float, but you’re returning a tuple.” That’s the entire feedback loop.

def train(loader, model):           # untyped — the checker can't help
    ...

def train(loader: DataLoader, model: Module) -> float:    # checker can verify callers
    ...

The cost is a few extra characters per signature. The benefit is the checker catching one real bug per kiloline of code, which compounds fast on a multi-author codebase.

Built-in scalars and containers

def normalise(x: float) -> float:
    return (x - mean) / std

def tokens_for(text: str) -> list[str]:
    return text.lower().split()

def vocab_from(tokens: list[str]) -> dict[str, int]:
    return {t: i for i, t in enumerate(set(tokens))}

def pair() -> tuple[int, str]:        # exactly two: int then str
    return 1, "x"

def coords() -> tuple[float, ...]:    # variable-length tuple of floats
    return (1.0, 2.0, 3.0)

list[int] (lowercase, 3.9+) is preferred over List[int] from the typing module. Same applies to dict, set, tuple, frozenset.

Optional, union, None

# Optional — value or None
def find_user(uid: int) -> User | None:
    ...

# Union — one of several types
def parse(value: str | bytes) -> dict:
    ...

# 3.9 and below would write Optional[User] / Union[str, bytes]

The X | Y syntax (PEP 604) is the modern form. Use it.

A common ML pattern: optional config fields.

def load_model(path: str, device: str | None = None) -> Module:
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    ...

Callable — function-as-argument

from typing import Callable

def train(
    loader: DataLoader,
    model: Module,
    loss_fn: Callable[[Tensor, Tensor], Tensor],   # takes (pred, target), returns loss
) -> float:
    ...

For full signature preservation when wrapping (e.g. decorators), use ParamSpec:

from typing import Callable, ParamSpec, TypeVar

P = ParamSpec("P")
R = TypeVar("R")

def timed(fn: Callable[P, R]) -> Callable[P, R]:
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        ...
        return fn(*args, **kwargs)
    return wrapper

This makes @timed preserve the wrapped function’s exact signature for the type checker.

TypeVar — generic functions and classes

When the exact type doesn’t matter but it must be the same in two places:

from typing import TypeVar

T = TypeVar("T")

def first(items: list[T]) -> T:
    return items[0]

>>> first([1, 2, 3])     # checker infers T = int, returns int
>>> first(["a", "b"])    # checker infers T = str, returns str

For generic classes:

from typing import Generic, TypeVar

T = TypeVar("T")

class Cache(Generic[T]):
    def __init__(self) -> None:
        self._data: dict[str, T] = {}

    def get(self, key: str) -> T | None:
        return self._data.get(key)

    def put(self, key: str, value: T) -> None:
        self._data[key] = value

embed_cache: Cache[Tensor] = Cache()

Python 3.12 added cleaner syntax under PEP 695:

class Cache[T]:
    def get(self, key: str) -> T | None: ...

— same meaning, less ceremony. Use it if your minimum Python is 3.12+.

Protocols — structural typing

The killer feature for ML codebases. A Protocol says “anything that has these methods is acceptable” — without requiring inheritance.

from typing import Protocol

class SupportsForward(Protocol):
    def forward(self, x: Tensor) -> Tensor: ...

def evaluate(model: SupportsForward, loader: DataLoader) -> float:
    for x, y in loader:
        pred = model.forward(x)
        ...

evaluate accepts anything with a forward(x: Tensor) -> Tensor method. No need to inherit from nn.Module, no isinstance check, no shared base class. This is duck typing made checkable.

This is what every “the function takes any iterable / any callable / any model” pattern wants. typing.Iterable, typing.Iterator, typing.Sequence, typing.Mapping are all built-in protocols you should reach for instead of concrete list / dict:

def total(values: Iterable[float]) -> float:    # accepts list, generator, set, tuple, ...
    return sum(values)

def total(values: list[float]) -> float:        # accepts ONLY list — too narrow
    return sum(values)

Default to Iterable[X] for “I just iterate”, Sequence[X] for “I also need indexing/len”, Mapping[K, V] for “key lookup”, and only narrow to list / dict when you mutate.

Literal — restrict to specific values

from typing import Literal

def reduce(values: Tensor, op: Literal["sum", "mean", "max"]) -> Tensor:
    if op == "sum":
        return values.sum()
    if op == "mean":
        return values.mean()
    return values.max()

The checker now knows op can only be one of three strings, so reduce(x, "sun") is a type error. Useful for the small “magic string” parameters every ML library has.

TypedDict — dicts with known keys

When you can’t / don’t want to use a dataclass but you have dicts with known structure (think: API responses, JSON configs):

from typing import TypedDict

class TrainResult(TypedDict):
    loss: float
    accuracy: float
    epoch: int

def train_one(loader, model) -> TrainResult:
    return {"loss": 0.3, "accuracy": 0.92, "epoch": 0}

Better: use a dataclass or Pydantic model. TypedDict is the right shape only when you’re stuck with a dict (e.g., a JSON-shaped API).

Annotated — type plus metadata

Used heavily by Pydantic and FastAPI:

from typing import Annotated
from pydantic import Field

class Config(BaseModel):
    learning_rate: Annotated[float, Field(gt=0, lt=1)]
    epochs: Annotated[int, Field(ge=1)]

The first argument is the type; subsequent ones are metadata that specific tools (Pydantic, here) interpret. Static type checkers ignore the metadata and treat Annotated[X, ...] as X.

Type checkers — pyright vs mypy

Both work. Use pyright if you start a new project today; it’s faster, the inference is better, and it’s what your editor is probably already running. Use mypy if your project standardizes on it.

Run them as part of CI:

# pyproject.toml
[tool.pyright]
include = ["src"]
strict = ["src/core"]

[tool.mypy]
python_version = "3.11"
strict = true

Common gotchas

• Type hints aren’t enforced at runtime. Pass a string where the hint says int and the function will run; you’ll just get a confusing error somewhere downstream. Use Pydantic if you need runtime enforcement.
• list[int] | None vs Optional[list[int]]. Same thing in 3.10+.
• Forward references for self-referential types: def f(self) -> "Tree": — string form, since Tree isn’t defined yet. Or from __future__ import annotations, which lazily evaluates all annotations.
• Any defeats the type system. Every Any is a hole. Sometimes necessary; reach for object (need to narrow before use) or Protocol (specify what you actually need) first.
• isinstance(x, list[int]) doesn’t work. Generics aren’t usable at runtime for instance checks. isinstance(x, list) does, but doesn’t check element types.
• Mutable defaults still bite. Type hints don’t fix def f(items: list[int] = []):. See PY 101.

A worked ML example — typed end-to-end

from typing import Iterable, Protocol, Literal, TypeVar
from dataclasses import dataclass
import torch
from torch import Tensor, nn

# A protocol that anything model-like satisfies
class Predictor(Protocol):
    def __call__(self, x: Tensor) -> Tensor: ...

# A typed result
@dataclass(frozen=True, slots=True)
class EvalResult:
    loss: float
    accuracy: float
    n_samples: int

# A typed function
def evaluate(
    model: Predictor,
    batches: Iterable[tuple[Tensor, Tensor]],
    *,
    reduction: Literal["mean", "sum"] = "mean",
) -> EvalResult:
    losses: list[float] = []
    correct = 0
    total = 0
    for x, y in batches:
        with torch.no_grad():
            pred = model(x)
        losses.append(nn.functional.cross_entropy(pred, y, reduction=reduction).item())
        correct += (pred.argmax(-1) == y).sum().item()
        total += y.numel()
    return EvalResult(
        loss=sum(losses) / len(losses),
        accuracy=correct / total,
        n_samples=total,
    )

Look at what the checker now enforces: any callable taking a Tensor and returning a Tensor is a valid model; the iterable can be a list, a DataLoader, or a generator; reduction can only be "mean" or "sum"; the return type is a fully-typed dataclass. None of this would crash at runtime if you got it wrong; the checker would catch it before you ran the code.

Where typing shows up in real ML codebases

• PyTorch has type stubs (torch.Tensor and friends are typed). Use the latest version for best stubs.
• Hugging Face Transformers is typed; pyright can autocomplete model.generate(...) arguments.
• Pydantic v2 uses annotations as the runtime spec for validation.
• FastAPI uses annotations to generate OpenAPI schemas and to validate request/response bodies.
• Hydra / OmegaConf use dataclasses (and their annotations) as the source of truth for configs.

The defensive habit: write the type hints as you write the function. Adding them later is always more painful, and you end up with a half-typed codebase that gives the checker fewer chances to help.

Resources

• Python docs — typing module — docs.python.org — the canonical reference.
• PEP 484 — Type Hints — peps.python.org — original spec.
• PEP 544 — Protocols (structural typing) — peps.python.org.
• PEP 612 — ParamSpec — peps.python.org — preserving signatures through decorators.
• PEP 695 — Type Parameter Syntax — peps.python.org — the cleaner generic syntax in 3.12+.
• mypy documentation — mypy.readthedocs.io — comprehensive but dense.
• Pyright documentation — github.com/microsoft/pyright — config reference and command-line usage.
• typeshed — github.com/python/typeshed — stubs for the standard library and third-party packages.