Skip to content

Generics & Protocols: Advanced Type Systems

Overview

Advanced typing enables sophisticated type systems: - Generic types: TypeVar, Generic for reusable components - Protocols: Structural typing without inheritance - Type checking: mypy validates complex types

TypeVar and Generics

TypeVar Basics

from typing import TypeVar, List

T = TypeVar('T')  # Unbounded type variable

def get_first(items: List[T]) -> T:
    """Returns first item - maintains type."""
    return items[0]

# Mypy infers types correctly
nums: List[int] = get_first([1, 2, 3])      # T = int
strs: List[str] = get_first(['a', 'b'])     # T = str

Bounded TypeVar

from typing import TypeVar, Union

Numeric = TypeVar('Numeric', int, float)  # Can be int or float

def add(a: Numeric, b: Numeric) -> Numeric:
    """Add two numeric values."""
    return a + b

# OK: both ints
result1 = add(1, 2)

# OK: both floats
result2 = add(1.0, 2.0)

# Error: mixed types
# result3 = add(1, 2.0)  # mypy error

Covariance and Contravariance

from typing import TypeVar

T_co = TypeVar('T_co', covariant=True)      # Output position
T_contra = TypeVar('T_contra', contravariant=True)  # Input position

class Producer(Generic[T_co]):
    def produce(self) -> T_co: ...

class Consumer(Generic[T_contra]):
    def consume(self, item: T_contra) -> None: ...

Protocols: Structural Subtyping

Protocol Definition

from typing import Protocol

class Drawable(Protocol):
    """Anything that can be drawn."""
    def draw(self) -> None: ...

class Circle:
    def draw(self) -> None:
        print("Drawing circle")

class Square:
    def draw(self) -> None:
        print("Drawing square")

def render(obj: Drawable) -> None:
    """Works with anything drawable."""
    obj.draw()

# Both work without inheriting from Drawable
render(Circle())   # OK
render(Square())   # OK

Runtime Checkable Protocols

from typing import Protocol, runtime_checkable

@runtime_checkable
class Comparable(Protocol):
    def __lt__(self, other) -> bool: ...

class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __lt__(self, other):
        return self.x + self.y < other.x + other.y

p1 = Point(1, 2)
p2 = Point(3, 4)

# Can check at runtime
if isinstance(p1, Comparable):
    print(f"p1 < p2: {p1 < p2}")

Real-World ML Typing

Model Interface Protocol

from typing import Protocol
import torch

class Trainable(Protocol):
    """Protocol for trainable models."""

    def forward(self, x: torch.Tensor) -> torch.Tensor: ...
    def parameters(self) -> list: ...

class SimpleModel:
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * 2

    def parameters(self) -> list:
        return [self.weight]

def train_model(model: Trainable, data: torch.Tensor) -> None:
    """Train any model conforming to Trainable."""
    output = model.forward(data)
    # Training code...

Typed Data Pipeline

from typing import Protocol, Iterator, Tuple

class DataLoader(Protocol):
    def __iter__(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]: ...

class SimpleDataLoader:
    def __iter__(self):
        for i in range(10):
            x = torch.randn(32, 10)
            y = torch.randint(0, 2, (32,))
            yield x, y

def train_with_loader(model, loader: DataLoader) -> None:
    """Train with any data loader."""
    for x, y in loader:
        # Training step...
        pass

Summary: Generics vs Protocols

Feature Generics Protocols
Reusability High High
Flexibility Constrained Flexible
Inheritance Required Not required
Type safety Strict Structural