writing

[Python] Typing the hard stuff

August 2023

Rules of Thumb

  1. Expect generic types, return specific types. This gives users of your functions and classes the biggest flexibility and safety
def process(items: Iterable[Item]) -> list[Item]:
    return [item for item in items if item.is_valid()]

How to type …

Classes

def birth(self, animal: type[Animal]) -> Animal:
    return animal()

Context Managers 1

from __future__ import annotations

from typing import overload
from types import TracebackType


class MyContextManager:
  def __enter__(self) -> None:
    pass

  @overload
  def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None:
    ...

  @overload
  def __exit__(
    self,
    exc_type: type[BaseException],
    exc_val: BaseException,
    exc_tb: TracebackType,
  ) -> None:
    ...

  def __exit__(
    self,
    exc_type: type[BaseException] | None,
    exc_val: BaseException | None,
    exc_tb: TracebackType | None,
  ) -> None:
    pass

Decorators

from typing import Callable, TypeVar, Any

T = TypeVar('T')

def foo(arg: str) -> Callable[[Callable[..., T]], Callable[..., T]]:
    def decorator(function: Callable[..., T]) -> Callable[..., T]:
        # Run on function registration only
        print(f"Decorator registered: {arg}")
        @functools.wraps(function)
        def wrapper(*args: Any, **kwargs: Any) -> T:
            # Run on every function call
            print(f"Decorated function call: {arg} - {args} {kwargs}")
            result = function(*args, **kwargs)
            return result
        # Note we can collapse the wrapper into just `function` if we don't need to modify the function call
        return wrapper
    return decorator

Optional Imports 1

try:
  import matplotlib.pyplot as plt

  _HAS_PLT = True
except ImportError:
    _HAS_PLT = False


def train():
  if _HAS_PLT:
    ...

External Modules 1

[tool.mypy]
mypy_path = "mypy_stubs"
mypy_stubs
└── example
    ├── __init__.pyi
    └── widgets.pyi
from typing import Any

def __getattr__(name: str) -> Any: ...

class Widget:
  def __init__(self, name: str) -> None: ...
  def frobnicate(self) -> None: ...
# In __init__.pyi
from typing import Any

def __getattr__(name: str) -> Any: ...

In runnable code, Python calls such a getattr function for any access of missing attributes. This allows you to do anything you like, such as dealing with complicated deprecations.

In stub files, type checkers understand this particular getattr definition to mark all unmentioned names as type Any. So, in our example, code that uses Widget can be fully type checked:

Duck-typed classes 1

from typing import Protocol

class SupportsClose(Protocol):
  def close(self) -> None: ...
# If you need arbitrary arguments
T = TypeVar("T", covariant=True)
Operation = Callable[..., T]
  1. From Adam Johnson’s Blog 2 3 4