Python Type Hints: How to Gradually Add Types for Third Party Packages

Gradually evolved cephalopods.

Hynek Schlawack recently described graduality as Python’s super power: the ability to prototype in the REPL, and gradually add linting, type checking, and other practices to refine your code into maintainable, production-ready software. You can also apply graduality within tools, activating checks one at a time and fixing the resulting errors as you go.

One place you can apply graduality with Mypy is in the type hints for third party packages. The default course of action is to add type hints for a whole package at once. You do this either through the package itself if it has a py.typed file, or with a *-stubs package (both specified in PEP 561). But when you add full hints for a package used widely in your code base, you might find yourself facing an insurmountable mountain of errors.

Instead, you can gradually add type hints for the package, piece by piece, fixing small batches of errors as you go. This iterative approach is more psychologically pleasing, and it reduces the chance for mistakes.

In this post, we’ll look at how to add such gradual type hints, using module-level __getattr__ functions in stub files, and a concrete use case from a Django project.

Bring your own stubs

Instead of using a package’s published type hints, you can provide your own stub files to Mypy. These can live on a directory specificed in the mypy_path option. For example, to use the directory mypy_stubs (in pyproject.toml):

[tool.mypy]
mypy_path = "mypy_stubs"

You can then put your .pyi stub files there, in directories mirroring the structure of the target packages. For example, if you were adding type hints for a package called example, with a widgets submodule:

mypy_stubs
└── example
    ├── __init__.pyi
    └── widgets.pyi

Okily dokily.

Use the module-level __getattr__

.pyi files contain bare function and class definitions with the type hints for the related code. All function bodies are marked empty with ....

For example, if you were specifying a Widget class (in example/widgets.pyi):

class Widget:
    def __init__(self, name: str) -> None: ...
    def frobnicate(self) -> None: ...

You can make a stub file “incomplete” by defining a module-level __getattr__ function that returns Any:

from typing import Any

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

class Widget:
    def __init__(self, name: str) -> None: ...
    def frobnicate(self) -> None: ...

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:

from example.widgets import Widget

widgets = [Widget("roundabout"), Widget("square")]

…but code using other names from the module uses Any, so type checking allows anything:

...

from example.widgets import deflagrate

deflagrate(widgets)  # could be an error!

Mypy does not look at the module’s code, so it can’t even check whether deflagrate() exists. Imports may even be typo’d, but hopefully you have tests to detect such basic bugs…

Anys all the way down

You can apply this __getattr__ pattern to add types for just part of a package. Imagine you wanted to have stubs only for example.widgets (for now). You’d layout the stub files like so:

mypy_stubs
└── example
    ├── __init__.pyi
    └── widgets.pyi

In the __init__.pyi, you only need the __getattr__:

from typing import Any

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

This will allow any other names in the package to be used.

Then you could fill in the type hints for widgets.pyi, perhaps completely for all names:

from collections.abc import Sequence

class Widget:
    def __init__(self, name: str) -> None: ...
    def frobnicate(self) -> None: ...

def deflagrate(widgets: Sequence[Widget]) -> None: ...

You can use such __init__.pyi files to an arbitrary depth:

mypy_stubs
└── example
    ├── __init__.pyi
    └── ham
        ├── __init__.pyi
        └── spam
            ├── __init__.pyi
            └── jam.pyi

This allows you to partially type a fraction of even quite large packages. Noice.

Gimme them Django enumeration types

Let’s look at a concrete example of applying this graduality.

I’ve been rolling out type hints at my client Silvr, gradually approaching Mypy’s “strict mode”. They have a big Django project, so at some point I’d like to add django-stubs. Unfortunately, adding all of django-stubs leads to thousands of errors… way too many for this grug-brained developer to tackle at once.

Without django-stubs installed, and the ignore_missing_imports option active, Mypy treats all types from django.* as Any. This meant no errors, except for one particular case.

Django provides enumeration types that subclass Python’s enum classes. Enums use some class definition magic to behave differently to normal classes, for which Mypy has special support.

Without type hints for Django, Mypy cannot tell that Django’s enumeration types are subclasses of enum.Enum. Thus, it does not extend the special behaviour to them. This minimal example exhibits the behaviour:

from django.db.models import IntegerChoices


class InvoiceStatus(IntegerChoices):
    UNPAID = 1, "Unpaid"
    PAID = 2, "Paid"
    DISINTEGRATED = 3, "Completely atomized"


class Invoice:
    ...


def find_invoices(status: InvoiceStatus) -> list[Invoice]:
    ...


find_invoices(InvoiceStatus.DISINTEGRATED)

At runtime, the enum class behaviour turns members like InvoiceStatus.DISINTEGRATED into an instance of InvoiceStatus. But without type hints for IntegerChoices, Mypy does not know this, and reports an error on the final line:

$ mypy example.py
example.py:18: error: Argument 1 to "find_invoices" has incompatible type "Tuple[int, str]"; expected "InvoiceStatus"
Found 1 error in 1 file (checked 1 source file)

Oh no.

Mypy sees IntegerChoices as only a normal class, with three attributes set to tuples.

Initially, we ignored such errors in Silvr by using tuple types:

def find_invoices(
    status: tuple[int, str],  # InvoiceStatus
) -> list[Invoice]:
    ...

…or with targetted type: ignore comments:

find_invoices(InvoiceStatus.DISINTEGRATED)  # type: ignore [arg-type]

Both techinques are unsatisfactory though, as they do not smoothly allow other parts of the Enum API to work, such as accessing the value attribute.

Partially importing django-stubs was the answer. I added the type stubs for just the enumeration types:

mypy_stubs/
└── django
    ├── __init__.pyi
    └── db
        ├── __init__.pyi
        └── models
            ├── __init__.pyi
            └── enums.pyi

django/db/models/enums.pyi came from django-stubs. The __init__.pyi modules all used the __getattr__ pattern as above, with one additional import in django/db/models/__init__.pyi:

from __future__ import annotations

from typing import Any

from .enums import Choices as Choices
from .enums import IntegerChoices as IntegerChoices
from .enums import TextChoices as TextChoices

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

With these files in place, no workarounds are needed to make full use of these types. Mypy is happy with the example file as-is:

$ mypy example.py
Success: no issues found in 1 source file

😊

I expect we’ll import more pieces of django-stubs into the Silvr project in order to approach Mypy’s strict mode. Then, finally when we make the switch to using the full django-stubs, the partial stub files can be deleted.

Fin

May you gradually refine your code forever,

—Adam


Learn how to make your tests run quickly in my book Speed Up Your Django Tests.


Subscribe via RSS, Twitter, or email:

One summary email a week, no spam, I pinky promise.

Related posts:

Tags: ,