Friday, 1 March 2019

Extending mypy with plugins

Python is a famously dynamic language. It’s easy to write DSL-like frameworks that are hard to understand by static type checkers. Although with recent mypy features, such as protocols and literal types, and with basic metaclass and descriptor support, we can more often provide precise types, it’s still often hard to avoid false positives and negatives. To solve this problem, while avoiding the need for custom type system features for every framework, mypy supports a plugin system. Plugins are Python modules that provide callbacks (plugin hooks) that mypy will call when type checking classes and functions that interact with a library or framework. For example, such hooks can provide a more precise function return type that is otherwise hard to express, or auto-generate some methods in a class, to reflect the effects of a class decorator. To read more about plugin system architecture and for the full list of available hooks, see these docs.

Bundled plugins for standard library

Mypy comes with a default plugin for some builtin functions and classes, and the ctypes, contextlib, and dataclasses modules. It also includes a plugin for attrs (for historical reasons — it was the first third-party plugin written for mypy). These plugins allow mypy to infer more precise types and correctly type check code using these library features. To illustrate this, consider this snippet: 

    from dataclasses import dataclass
    from typing import Generic, TypeVar
    
    @dataclass
    class TaggedVector(Generic[T]):
        data: List[T]
        tag: str
    
    position = TaggedVector([0, 0, 0], 'origin')

Above, get_class_decorator_hook() is called when analyzing the class definition. It adds autogenerated methods, including __init__(), to the class body. Mypy uses this generated constructor to correctly infer TaggedVector[int] as the type of position. As you can see, plugins work even with generic classes.

Here’s another snippet: 

    from contextlib import contextmanager
    
    @contextmanager
    def timer(title: str) -> Iterator[float]:
        ...
    with timer(9000) as tm:
        ...

The get_function_hook() hook provides a precise return type for the contextmanager decorator, so that calls to the decorated function can be type-checked precisely. Mypy can now detect an error: the argument to timer() must be a string.

Combining plugins and stubs

In addition to relying on dynamic Python features, frameworks often have the challenge of having large APIs. Mypy needs stub files for the libraries to check code that uses those libraries (unless the library contains inline annotations, which is still often not the case). Distributing the stubs for large frameworks via typeshed is not very practical:

  • Typeshed has a relatively slow release cycle (it’s bundled with mypy).
  • Incomplete stubs can cause false positives that are hard to avoid.
  • There’s no easy way to mix and match stubs from different typeshed versions.
Stub packages, introduced in PEP 561, help with these issues:
  • Maintainers can release a stub package as often as they want.
  • Users who haven’t opted in to using the package won’t see any false positives.
  • You can freely install arbitrary versions of multiple different stub packages.

Moreover, pip allows combining stubs for a library and the corresponding mypy plugin into a single distribution. Stubs for a framework and a corresponding mypy plugin can now be easily developed and distributed together, which is often useful since plugins fill in missing or imprecise definitions in the stubs.

A recent example of such a package is SQLAlchemy stubs and plugin, the first public 0.1 release of which was published on PyPI earlier this month. Although the project is still in early alpha, we are already successfully using it at Dropbox to improve type checking precision. Currently the plugin understands basic ORM declarations: 

    from sqlalchemy.ext.declarative import declarative_base
    from sqlalchemy import Column, Integer, String
    
    Base = declarative_base()
    
    class User(Base):
        __tablename__ = 'users'
        id = Column(Integer, primary_key=True)
        name = Column(String)

In the above snippet, the plugin uses the get_dynamic_class_hook() hook to tell mypy that Base is a valid base class, even though it doesn’t look like one. Then get_base_class_hook() is called on the definition of User, and it adds some autogenerated attributes. Next we instantiate the model: 

    user = User(id=42, name=42)

The get_function_hook() hook is called, so that mypy can spot the error here: an integer is given instead of a user name.

The stubs define Column as a generic descriptor, so that attributes on a model get correct types: 

    id_col = User.id  # Inferred type is "Column[int]"
    name = user.name  # Inferred type is "Optional[str]"

We welcome PRs that add more precise types to the stubs (the progress for core modules is tracked here).

Here are some gotchas that we found while working on the stubs:

  • Use a module-level __getattr__() to avoid false positives during the very early stages, when stubs are incomplete (this suppresses mypy errors due to missing module attributes). You can also use it in __init__.py files if some submodules are missing.
  • Descriptors often help with giving more precise types for customized attribute access (like in the Column example above), and it’s OK to use them even if the actual runtime implementation uses a more complex mechanism involving a metaclass, for example.
  • Don’t hesitate to declare framework classes as generic in stubs. Although they are not generic at runtime, this often allows us to give much more precise types for certain framework features, while runtime errors can easily be worked around. (We hope that frameworks will gradually add built-in support for generic types by explicitly inheriting relevant classes from typing.Generic.)

Recently released mypy plugins

There are already several plugins available for some popular Python frameworks. Apart from the above mentioned plugin for SQLAlchemy, other notable recent examples of packages with stubs and a bundled mypy plugin include stubs for Django and Zope Interfaces. All these projects are under active development.

Installing and enabling plugin + stub packages

Use pip to install a package with a mypy plugin and/or stubs in the virtual environment where mypy is installed: 

    $ pip install sqlalchemy-stubs

Mypy should automatically discover the installed stubs. To enable the plugins you’ve installed, explicitly include them in your mypy.ini (or a custom config file): 

    [mypy]
    plugins = sqlmypy, mypy_django_plugin.main

Developing mypy plugins and writing stubs

If you want to develop a stubs + plugin package for a framework you use, you can use the sqlalchemy-stubs repository as a template. It includes the setup.py file, testing infrastructure using data-driven tests, and an example plugin class with a bunch of plugin hooks. We recommend using stubgen, an automatic stub generator that comes with mypy, to get started with stubs. Stubgen got several improvements in mypy 0.670.

For more details about mypy plugin system, see the docs. You can also just browse the source code of the plugins mentioned above. If you have any questions, feel free to ask them on the Python typing Gitter chat.

Posted by at