samgeo.codes

Musings on Python Type Hints

"Gradual Typing" has become incredibly popular over the last 8 years or so. The most notable examples of this phenomena exist in the JavaScript space. Since JavaScript is the lingua franca of the web there have been several efforts to leverage the benefits of static type systems to enable easier programming-in-the-large for JavaScript. Gradual type systems such as TypeScript and Flow have found the most success in this space.

Shortly after the release of TypeScript, Guido van Rossum and Ivan Levkivskyi created PEP 483 proposing a type hinting system for Python. Notably absent from this proposal is any "official" type checking program for verifying the correctness of type hints. As a result, there are four "major" type checkers for python:

  • mypy: A very early project (starting in 2012) from Jukka Lehtosalo. mypy and its predecessors have heavily influenced direction of python type hints. mypy has been supported heavily by Jukka's employer Dropbox. mypy is arguably the most popular type checker for Python.
  • pyright: A type checker built by Microsoft with great VS Code and LSP integration. Pyright was first released in 2019 and is bundled in the Microsoft Python extension for VS Code.
  • pytype: A static type analyzer built by Google. Pytype was first released in 2018 and emphasizes type inference and local leniency.
  • pyre: A static type checker built by Meta. Designed for large code bases and first-class integration with Buck, Meta's build system. The first release I can find dates back to 2019.

There have been countless arguments over the value of static types over the years. I really like having a type system that can help me to avoid some common mistakes and guide me when refactoring. Most of all, however, type hints should broadcast intent; they are one of the main forms of documentation that the computer can help to verify.

Since starting a new job earlier this year, I've spent a good amount of time learning how to leverage type hints to help write readable, maintainable code. Let's take a look at some patterns that arise in Python where type hinting can get a little bit tricky.

Getting started

I'll be using mypy since it seems to be the most widely used tool, at the moment.

If you don't already have python installed, follow the instructions on the official python website to install python for your platform. At the time of writing, the latest stable version is 3.10.7 which is what I will be using for my examples.

With those out of the way, we can start our python project and install mypy:

$ mkdir python_types && cd python_types
$ python3 -m venv venv
$ source venv/bin/activate  # may be different on Windows
$ python3 -m pip install mypy

"Hello, world!" of type checking

Let's start with a simple example:

def add(x, y):
    return x + y


print(add(1, 2))

We can run this:

$ python3 -m src.hello
3

Now, let's see what mypy has to say:

$ mypy --strict src/hello.py
src/hello.py:1: error: Function is missing a type annotation
src/hello.py:5: error: Call to untyped function "add" in typed context
Found 2 errors in 1 file (checked 1 source file)

In strict mode, mypy expects type annotations in certain locations. In particular, function parameters and return types must be annotated.

Let's modify our program and try again:

def add(x: int, y: int) -> int:
    return x + y


print(add(1, 2))

When we run mypy now:

$ mypy
Success: no issues found in 1 source file

Cool.

Now you may have noticed that there are some seemingly valid uses of the add function that mypy now rejects; for example:

def add(x: int, y: int) -> int:
    return x + y


print(add("foo", "bar"))

This program runs just fine:

$ python3 -m src.hello
foobar

But when we run mypy:

$ mypy
src/hello.py:5: error: Argument 1 to "add" has incompatible type "str"; expected "int"
src/hello.py:5: error: Argument 2 to "add" has incompatible type "str"; expected "int"
Found 2 errors in 1 file (checked 1 source file)

Of course, if we annotate the function parameters as being integers, any reasonable type checker should complain here. But a lot of the power of using python comes from being able to use functions in a flexible manner.

Fortunately, there are features in the standard library that enable us to create better annotations fairly easily.

Let's start with TypeVar (introduced in pep 483). If you are familiar with generics in another language such as TypeScript or Java, TypeVar enables something similar.

Let's take a look at what our function signature might look like with the use of TypeVar:

from typing import TypeVar


T = TypeVar("T")


def add(x: T, y: T) -> T:
    return x + y


print(add("foo", "bar"))

This looks close to reasonable, let's make sure it runs:

$ python3 -m src.hello
foobar

Nice. Let's see if mypy likes it:

$ mypy
src/hello.py:8: error: Returning Any from function declared to return "T"
src/hello.py:8: error: Unsupported left operand type for + ("T")
Found 2 errors in 1 file (checked 1 source file)

Huh. Well, when you think about it, just because x and y have the same type doesn't necessarily mean that x and y can be added together. Maybe x and y are dictionaries; the add operation is not defined for dictionaries.

Perhaps we can do a little bit better. The critical idea here is called a Protocol (introduced in pep 544). A Protocol lets us specify the shape of objects that we expect to see. In the literature, you might see this described as "structural subtyping" but most of the time you'll hear it as "duck typing": if it looks like a duck and quacks like a duck, it must be a duck.

Let's see if we can get mypy to approve our program:

from typing import Protocol, TypeVar

Self = TypeVar("Self", bound="Addable")

class Addable(Protocol):
    def __add__(self: Self, other: Self) -> Self:
        ...


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


def add(x: T, y: T) -> T:
    return x + y


print(add("foo", "bar"))

Whew, that's quite a bit of typing for such a simple example.

Let's at least see if it works. First we run it:

$ python3 -m src.hello
foobar

Now let's run mypy:

$ mypy src/hello.py
Success: no issues found in 1 source file

Hooray! This is a lot of machinery for such a simple example. Let's unpack what's going on here:

Self = TypeVar("Self", bound="Addable")

class Addable(Protocol):
    def __add__(self: Self, other: Self) -> Self:
        ...

This allows us to refer to the type of objects that can be added to other objects of the same type to produce yet another object of that same type.

The Self nonsense going on here allows us to constrain the type of the value that can be added to only those of the same type (and constrain the return type to the same type). This gets a lot easier in python 3.11 with the introduction of the Self type to the standard library (pep 673).

Anthony Sotille has a great video covering the Self type, as well: link. As an aside, many of Anthony's videos are great bite-sized examples of various python concepts; he also has a great live stream!

Now, this protocol isn't perfect: addition may be defined for combinations of types. As far as I know, there isn't a great way to express that sort of thing, just yet. Furthermore, while mypy accepts this program, both pyre and pyright seem to reject it. This is one of the downsides of having several different type checker implementations. Hopefully these discrepancies get ironed out as the ecosystem matures.

With that introduction out of the way, let's take a look at a pattern that seem to come up pretty often in python and how we might be able to add type annotations.

Decorator factories

Decorators are a curiously functional feature of python. Decorator factories (functions that produce decorators) are fairly common but surprisingly difficult to write type annotations for.

We're going to take a look at a much more complex example than the add case that we explored above. In particular, we are going to leverage asynchronous code! If you aren't familiar with asyncio in python, I encourage you to explore the official documentation: asyncio docs

To start out, we have a complete implementation of a decorator that continuously prints out the running time of the decorated async function:

import asyncio
import functools
import inspect
import time
from typing import Final


async def print_with_timer(format: str, interval: float = 0.1) -> None:
    start = time.time()
    while True:
        elapsed = time.time() - start
        print(format.format(elapsed=elapsed), end="")
        await asyncio.sleep(interval)


OVERWRITE_LINE: Final[str] = "\033[F"
TIMER_FORMAT: Final[str] = f"\n{OVERWRITE_LINE}Invocation {{module}}.{{name}}({{arguments}}) has been running for {{{{elapsed:0.2f}}}}s"


def with_timer(interval_or_callback):
    def inner(func):
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            arguments = inspect.signature(func).bind(*args, **kwargs).arguments
            argument_string = ", ".join(f"{name} = {value!r}" for name, value in arguments.items())
            timer_format = TIMER_FORMAT.format(module=func.__module__, name=func.__qualname__, arguments=argument_string)
            if isinstance(interval_or_callback, float):
                task = asyncio.create_task(print_with_timer(timer_format, interval_or_callback))
            else:
                task = asyncio.create_task(print_with_timer(timer_format))
            try:
                result = await func(*args, **kwargs)
            finally:
                task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
            print()
            return result
        return wrapper

    if isinstance(interval_or_callback, float):
        return inner

    return inner(interval_or_callback)


@with_timer
async def do_stuff() -> int:
    await asyncio.sleep(5.0)
    return 42


@with_timer(0.5)
async def do_more_stuff(duration: float) -> str:
    await asyncio.sleep(duration)
    return "Do a good turn daily."


async def main() -> int:
    important_number = await do_stuff()
    print(f"The meaning of life is {important_number}")
    scout_slogan = await do_more_stuff(3.0)
    print(f"The scout slogan is \"{scout_slogan}\"")
    return 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))

First, let's run it to see what this program does:

$ python3 -m src.timer
Invocation __main__.do_stuff() has been running for 4.97s
The meaning of life is 42
Invocation __main__.do_more_stuff(duration = 3.0) has been running for 2.51s
The scout slogan is "Do a good turn daily."

That's neat! We can see which coroutine we are waiting for and how long it has been running.

Let's see if mypy likes this:

$ mypy --strict src/untyped_timer.py
src/untyped_timer.py:20: error: Function is missing a type annotation
src/untyped_timer.py:21: error: Function is missing a type annotation
src/untyped_timer.py:23: error: Function is missing a type annotation
src/untyped_timer.py:46: error: Call to untyped function "inner" in typed context
src/untyped_timer.py:49: error: Untyped decorator makes function "do_stuff" untyped
src/untyped_timer.py:55: error: Call to untyped function "with_timer" in typed context
src/untyped_timer.py:55: error: Untyped decorator makes function "do_more_stuff" untyped
Found 7 errors in 1 file (checked 1 source file)

Aha, mypy loses track of the types of do_stuff and do_more_stuff since we haven't added type hints to the decorator.

This isn't a huge deal, here, but if we are relying on mypy to be able to determine the type of an intermediate expression as we are coding, this can be a pretty big pain. For example, let's try revealing the type of do_stuff:

...
reveal_type(do_stuff)
reveal_type(do_more_stuff)
...

When we run mypy:

$ mypy --strict src/untyped_timer.py
src/untyped_timer.py:20: error: Function is missing a type annotation
src/untyped_timer.py:21: error: Function is missing a type annotation
src/untyped_timer.py:23: error: Function is missing a type annotation
src/untyped_timer.py:46: error: Call to untyped function "inner" in typed context
src/untyped_timer.py:49: error: Untyped decorator makes function "do_stuff" untyped
src/untyped_timer.py:55: error: Call to untyped function "with_timer" in typed context
src/untyped_timer.py:55: error: Untyped decorator makes function "do_more_stuff" untyped
src/untyped_timer.py:61: note: Revealed type is "Any"
src/untyped_timer.py:62: note: Revealed type is "Any"
Found 7 errors in 1 file (checked 1 source file)

We see that the revealed type is Any which isn't very helpful.

If we remove the decorator, let's see if we can get a little bit more help:

$ mypy --strict src/untyped_timer.py
src/untyped_timer.py:20: error: Function is missing a type annotation
src/untyped_timer.py:21: error: Function is missing a type annotation
src/untyped_timer.py:23: error: Function is missing a type annotation
src/untyped_timer.py:46: error: Call to untyped function "inner" in typed context
src/untyped_timer.py:59: note: Revealed type is "def () -> typing.Coroutine[Any, Any, builtins.int]"
src/untyped_timer.py:60: note: Revealed type is "def (duration: builtins.float) -> typing.Coroutine[Any, Any, builtins.str]"
Found 4 errors in 1 file (checked 1 source file)

So we see that mypy can handle this sort of thing, in general. How can we leverage type hints to give mypy enough information to tell us the type with the decorator in use?

Let's dive in!

import asyncio
import functools
import inspect
import time
from typing import Awaitable, Callable, Final, overload, ParamSpec, Protocol, TypeVar


async def print_with_timer(format: str, interval: float = 0.1) -> None:
    start = time.time()
    while True:
        elapsed = time.time() - start
        print(format.format(elapsed=elapsed), end="")
        await asyncio.sleep(interval)


OVERWRITE_LINE: Final[str] = "\033[F"
TIMER_FORMAT: Final[str] = f"\n{OVERWRITE_LINE}Invocation {{module}}.{{name}}({{arguments}}) has been running for {{{{elapsed:0.2f}}}}s"


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

class Decorator(Protocol):
    def __call__(self, func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
        ...


@overload
def with_timer(interval_or_callback: float) -> Decorator:
    ...

@overload
def with_timer(interval_or_callback: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
    ...

def with_timer(interval_or_callback: float | Callable[P, Awaitable[R]]) -> Decorator | Callable[P, Awaitable[R]]:
    def inner(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
        @functools.wraps(func)
        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            arguments = inspect.signature(func).bind(*args, **kwargs).arguments
            argument_string = ", ".join(f"{name} = {value!r}" for name, value in arguments.items())
            timer_format = TIMER_FORMAT.format(module=func.__module__, name=func.__qualname__, arguments=argument_string)
            if isinstance(interval_or_callback, float):
                task = asyncio.create_task(print_with_timer(timer_format, interval_or_callback))
            else:
                task = asyncio.create_task(print_with_timer(timer_format))
            try:
                result = await func(*args, **kwargs)
            finally:
                task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
            print()
            return result
        return wrapper

    if isinstance(interval_or_callback, float):
        return inner

    return inner(interval_or_callback)


@with_timer
async def do_stuff() -> int:
    await asyncio.sleep(5.0)
    return 42


@with_timer(0.5)
async def do_more_stuff(duration: float) -> str:
    await asyncio.sleep(duration)
    return "Do a good turn daily."


async def main() -> int:
    important_number = await do_stuff()
    print(f"The meaning of life is {important_number}")
    scout_slogan = await do_more_stuff(3.0)
    print(f"The scout slogan is \"{scout_slogan}\"")
    return 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))

There's a lot going on here, so let's walk through some of the new additions:

  • ParamSpec: this enables us to capture the signature of a Callable object. You can do some pretty powerful things including adding additional arguments to functions while enabling your type checker to still understand the types being used. See pep 612 to read more about how ParamSpec works.
  • Decorator protocol: in many cases, mypy doesn't actually force us to use a protocol for this sort of thing, but it is actually pretty nice. This basically says that our decorator can take any callable object that takes in arguments that adhere to parameter specification P and returns objects of type R and our decorator will return a callable object with that same type.
  • Overload: we want users to be able to use with_timer with and without an interval. Always having to write with_timer() is kind of awkward, so we can branch on the type that we receive as input. overload informs the type checker that the return type of our decorator factory is always consistent depending on the input type to our decorator. Details on overload can be found in pep 484.

By leveraging these ideas, we are able to get mypy to fully understand what is going on with the types of our decorated functions.

Let's first run mypy and see if it works:

$ mypy --strict src/timer.py
Success: no issues found in 1 source file

Great!

Let's see if mypy can correctly identify the types of our decorated functions by adding reveal_type

...
reveal_type(do_stuff)
reveal_type(do_more_stuff)
...

Now we run mypy again:

$ mypy --strict src/timer.py
src/timer.py:75: note: Revealed type is "def () -> typing.Awaitable[builtins.int]"
src/timer.py:76: note: Revealed type is "def (duration: builtins.float) -> typing.Awaitable[builtins.str]"
Success: no issues found in 1 source file

Pretty neat!

Wrapping up

You can see that as library authors, there are lots of tools at our disposal for ensuring that we can produce well-typed APIs for our users while still leveraging many of the dynamic features of python.

There are plenty more things to learn about in the world of python type hints. If these ideas are exciting to you, dive into the peps and consider contributing to one of the type checking projects!

The static typing story for python is still very young. We don't have many of the type computation facilities of TypeScript. But the story is getting better! And from my perspective, having another way to help communicate intent and enable users to better understand interfaces is fantastic.