Search code examples
pythondjangodjango-ormpython-typingpyright

How to annotate type of Manager().from_queryset()?


In Django I have custom QuerySet and Manager:

from django.db import models

class CustomQuerySet(models.QuerySet):
    def live(self):
        return self.filter(is_draft=False)

class CustomManager(models.Manager):
    def publish(self, instance: "MyModel"):
        instance.is_draft = False
        instance.save()

In my model I want to use both, so I use from_queryset method:

class MyModel(models.Model):
    objects: CustomManager = CustomManager().from_queryset(CustomQuerySet)()

    is_draft = models.BooleanField(blank=True, default=True)

Since I annotated objects as CustomManager, Pylance (via vscode) logically yells at me that MyModel.objects.live() is wrong, due to Cannot access attribute "live" for class "CustomManager" Attribute "live" is unknown.

Removing type annotation leads to similiar complaint: Cannot access attribute "live" for class "BaseManager[MyModel]" Attribute "live" is unknown.

How to annotate objects in MyModel so Pylance will be aware that objects also has CustomQuerySet methods available, not only CustomManager methods?


Looking at Django's source from_queryset constructs a new subclass of CustomManager by iterating through CustomQuerySet methods:

@classmethod
def from_queryset(cls, queryset_class, class_name=None):
    if class_name is None:
        class_name = "%sFrom%s" % (cls.__name__, queryset_class.__name__)
    return type(
        class_name,
        (cls,),
        {
            "_queryset_class": queryset_class,
            **cls._get_queryset_methods(queryset_class),
        },
    )

So as @chepner pointed out in his comment we get a structural subtype of CustomManager, whose _queryset_class attribute is CustomQuerySet. So the question fundamentally is: how to type annotate that dynamically generated subclass in a way at least good enough for type-checker and autocomplete to work?

Approaches that I looked at so far are unsatisficatory:

  1. Django doesn't type annotate return of from_queryset.
  2. Django-types annotates it as type[Self], which lacks CustomQuerySet methods.
  3. Same with Django-stubs.

Solution

  • CustomManager.from_queryset(CustomQuerySet) creates a new subclass of CustomManager at runtime that copies a set of attributes from CustomQuerySet. As such, the "correct" solution would be to define a protocol that enumerates the attributes copied by the private Manager._get_queryset_methods used by from_queryset.

    As that is likely a lot of work (and fragile, in that it relies on private implementation details that may change), the broader quick-and-dirty method you propose of pretending that the object is an instance of both CustomManager and CustomQuerySet may be sufficient for your needs.

    class ManagerQuerySet(CustomManager, CustomQuerySet):
        pass
    
    
    class MyModel(models.Model):
        objects: ManagerQuerySet = CustomManager().from_queryset(CustomQuerySet)()
    
        is_draft = models.BooleanField(blank=True, default=True)
    

    (Really, from_queryset seems to perform a kind of "structural" inheritance, which may not be complete, but good enough to permit the white lie that the new class is a nominal subclass of CustomQuerySet.)