CODEWITHSUNDEEP
pythondjangodjango-ormpython-typingpyright
Matija Sirk
Matija Sirk

Reputation: 960

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.

Upvotes: 1

Views: 83

Answers (2)

chepner
chepner

Reputation: 532093

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.)

Upvotes: 1

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 477523

It does not seem to make much sense to place the publish(..) at the manager level, since one might want to publish all sorts of QuerySets.

You typically define this on the QuerySet with:

class CustomQuerySet(models.QuerySet):
    def live(self) -> CustomQuerySet['MyModel']:
        return self.filter(is_draft=False)

    def publish(self):
        return self.update(is_draft=False)

and then you inject the CustomQuerySet with:

class MyModel(models.Model):
    objects: QuerySet['MyModel'] = CustomQuerySet.as_manager()

You can then publish a set of items with the QuerySet, like:

MyModel.objects.filter(pk=1).publish()

Upvotes: 0

Related Questions