How to hint the type of Django's model field `objects` to a dynamically generated class?

Viewed 2546

I have a Team model in my Django project. I create its custom model manager with QuerySet.as_manager().

class TeamQuerySet(models.QuerySet):
    def active(self) -> "models.QuerySet[Team]":
        return self.filter(is_active=True)

class Team(models.Model):
    is_active = models.BooleanField()

    objects = TeamQuerySet.as_manager()

When I try to execute Team.objects.active(), mypy gives the following error:

error: "Manager[Any]" has no attribute "active"
In [5]: Team.objects
Out[5]: <django.db.models.manager.ManagerFromTeamQuerySet at 0x10eee1f70>

If I was explicitly defining a TeamManager class, there would be not a problem. How can I hint the type of Django model field objects to a dynamically generated class?

3 Answers

Based on the Manager[Any] I assume you are already using django-stubs.

Unfortunately it looks like there's an open issue to make QuerySet.as_manager generic over the model it's attached to that has not been resolved yet.

Even if the PR addressing the issue got merged I'm afraid it wouldn't address your immediate issue because the as_manager needs to be generic over the generic QuerySet subclass used to create the manager in order for both .active to be available and attributes relating to Team be available.

In this regard this other PR, which is unfortunately quite stale, seems to properly address your issue.

I've worked around this with a little switch-a-roo for MyPy's sake:

_Q = TypeVar("_Q", bound="WorkflowQuerySet")


class WorkflowQuerySet(models.QuerySet["WorkflowModel"]):
    """
    Queryset for workflow objects.
    """

    def count_objects(self) -> int:
        raise NotImplementedError

    def latest_objects(self: _Q) -> _Q:
        raise NotImplementedError


if TYPE_CHECKING:
    # Create a type MyPy understands
    class WorkflowManager(models.Manager["WorkflowModel"]):
        def count_objects(self) -> int:
            ...

        def latest_objects(self) -> _Q:
            ...


else:
    WorkflowManager = WorkflowQuerySet.as_manager


class WorkflowModel(models.Model):
    """
    A model that has workflow.
    """

    objects = WorkflowManager()

Here is my answer using generics and typevar

from typing import Generic, TypeVar
from django.db import models

class BookQueryset(models.QuerySet['Book']):
    ...

class Book(models.Model):
   objects: BookQueryset = BookQueryset.as_manager()


book = Book.objects.all()[0]

If you inspect book is type Book

Related