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?

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

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

Back to Top