Introducing the Concrete[] and DefaultQuerySet[] annotations

[delfick]

Hi!

In mypy 1.5.1 mypy/django-stubs started understanding that abstract models don't have a .objects and for the 6 million line django monolith I work on this caused us a lot of errors, so I made a fix for that, which also manages to fix a wider issue.

I've been waiting till we finally finished getting to latest mypy/django-stubs before I start a discussion about it. It took quite a while, but finally we're there. It does a lot so I'm unsure how much can be taken into this project, but nonetheless:

https://github.com/delfick/extended-mypy-django-plugin

The problem

When we have a variable whose type is a specific Abstract Django ORM model, that variable statically doesn't have the attributes/methods on it that appear on all the concrete modles underneath it that aren't also on the Abstract model.

Adding type annotations to the abstract model for these attributes/methods can't always work because those annotations won't themselves get attributes/methods that are dynamically discovered by django-stubs.

My extension provides a way of dealing with this limitation.

The solution

In the codebase I work on, it is very common to want to operate on a model that we know is a concrete variant of some specific abstract model, without knowing the specific models that this instance could be. (Our monolith has a lot of different configurations to support installations all over the world).

def do_something(my_model: SomeAbstractModel) -> str:
    # This gets a [attr-defined] error cause abstract models don't have reverse relationships on them
    return my_model.some_reverse_relationship.latest().some_property

By adding Concrete to the annotation from the plugin, we can now get the attributes correctly

def do_something(my_model: Concrete[SomeAbstractModel]) -> str:
    # no attr-defined error because my_model is the union of only the concrete models
    return my_model.some_reverse_relationship.latest().some_property

The plugin uses a mypy hook to see where this annotation is used and replace it with a union of all of the concrete classes that inherit from the abstract model. Because a union in mypy is a set intersection, not including the abstract model in that union means we can still see the attributes/methods common to the concrete models that aren't on the abstract model.

Other functionality available

Concrete.cast_as_concrete does a cast of something typed as an abstract to make mypy think it's type is a union of all the Concrete.

And DefaultQuerySet finds the type of the default queryset. So for a model that doesn't have a custom queryset it'd be django.db.models.QuerySet[MyModel] whereas for a model with a custom objects manager that has a custom queryset, it'd find that queryset. Using the annotation on an abstract model finds the union of these querysets.

More information on can be found in the usage documentation: https://extended-mypy-django-plugin.readthedocs.io/en/latest/api/usage.html

Limitation: Type variables

The main limitation the plugin has is that we can't use the annotations on type variables. This was originally supported but I had to remove that capability because I can't do enough work during the semantic phase for it to not cause problems.

The good news is that, at least in our codebase, that's not a big loss. For example the main case was needing to do a Concrete[Self] but we can get away with returning typing.Self and using Concrete.cast_as_concrete inside the method:

class MyAbstractModel(models.Model):
    class Meta:
        abstract = True

    def new(self) -> Self:
        concrete = Concrete.cast_as_concrete(self)
        # concrete is typed as the union of the concrete models
        # rather than as an instance of this abstract model
        return concrete.objects.create()

This is enough because it's the caller that determines the specific model being operated on, and in the method, cast_as_concrete can be used to get a type that represents all the known concrete models for this abstract model.

The complexity of how it works

The plugin works well, but has to deal with some unavoidable complexity.

Incremental cache and virtual dependencies

Mypy plugins can't directly make mypy re-analyze specific files, and they're unable to directly edit the incremental cache. So we need a way to resolve all subclasses of any abstract models before we attempt to type check. Like django-stubs, I use django introspection to resolve those subclasses (and custom querysets), however, INSTALLED_APPS can change what concrete models are available with no changes to the code itself. I've solved this by introducing the concept of "virtual dependencies".

In the plugin, we generate "virtual dependencies" and return them as dependencies during the get_additional_deps hook. This way when the django ORM models change, I can change the relevant virtual dependencies, which causes mypy to re-analyze the files as needed. This includes models that aren't part of an installed app - they exist and are analyzed but aren't seen by Django introspection. My tests have an example of such a virtual dependency.

When it comes to replacing the Concrete[] and DefaultQuerySet[] annotations themselves, I can use those virtual dependencies to make this really quite simple. The virtual dependencies are themselves python files and I'm able to fill them with the union types such that I can replace the annotations with a reference to the appropriate union type.

Mypy daemon

When the INSTALLED_APPS setting changes the django introspection needs to be redone in a new process because Django is built on global registries and cannot simply be reloaded in the same process. This is a problem for the mypy daemon (and is an existing problem with django-stubs as well). To make this work I figured out that I can force dmypy to restart by relying on how the dmypy daemon will discover a "version" property next to where it found the plugin. The dmypy daemon will then restart if that string has changed. I was able to make this "version" string dynamic by relying on locals(). This involved making it so that when you define the plugin entry point, you define this as an instance of a special object as opposed to a simple function.

That special object ensures that when mypy reaches for the plugin, django introspection will be done in a new process, figure out what the new version string is, then the plugin will be returned for mypy to use. Ideally I could do this in an external process that the plugin controls, but dmypy would still need to be restarted because this is an existing problem with django-stubs itself. Fixing this would require heavy surgery to django-stubs because django objects aren't serializeable between processes.

When I run the django introspection in an external process, it's running this script. This works by looking at all the enabled plugins until it finds the one that has the object that replaces the plugin entry point as this object is given the knowledge of how to create the virtual dependencies.

As a nice bonus, this means that we can customize how to create the django environment when we define the plugin entry point. I have an example project of such customization in the docs showing how it looks with something like django-configurations.

Testing

I'm a strong believer in tests and to make all of this reliable does require tests. I was able to use pytest-mypy-plugins as a base to make my own plugin for testing type checkers and I'm able to now write tests in python such that I can test what happens when mypy runs after changing the code.

Plugin

I also needed the ability to write this plugin as an extension of django-stubs given that the mypy plugin system is written such that first plugin to respond to a hook wins. To make this less awful I wrote a layer that makes it a bit easier to separate choosing the hook from the action of the hook. so the plugin looks like this.

Moving forward

I'd love to hear thoughts from django-stubs maintainers of how objectional what I have created is here, or if this is something you'd be interested in bringing into django-stubs itself in some fashion. I also feel quite isolated from django-stubs/mypy development and I imagine changes to the mypy plugin system could remove some of the more crafty things I do to make this work. But I'm unsure how to go about creating that kind of movement.