Type hint custom Django model manager - correct way?

[dfillingham]

Hi,

What is the correct way to type hint a custom django model manager so that Intellisense will work when interacting with the manager via .objects and also when returning model instances from the built-in django queryset APIs?

Python version - 3.10.12
Django version - 5.0
Pylance version - v2024.6.1

A simple model & manager to demonstrate the issue I'm facing:

from __future__ import annotations

import typing
from django.db import models
from django.db.models import Manager

if typing.TYPE_CHECKING:
    from typing_extensions import Self

class SalesOrderManager(models.Manager):
    
    def create_sales_order(self, order_number: str):
        self.create(order_number=order_number, status="initial-entry")


class SalesOrder(models.Model):
    order_number = models.CharField()
    status = models.CharField()

    objects = SalesOrderManager()

    def is_order_open(self) -> bool:
        open_statuses = ["initial-entry", "awaiting-payment", "awaiting-delivery"]

        return self.status in open_statuses


def some_func():
    SalesOrder.objects.create_sales_order(order_number="123")

    order = SalesOrder.objects.first()

    is_open = order.is_order_open()

If I annotate objects like so: objects: SalesOrderManager = SalesOrderManager() then Intellisense works to find the custom function create_sales_order on the manager, however built-in Django APIs such as get, first etc are not typed as SalesOrder so the is_order_open function isn't found, nor are the model's fields.

This also occurs if I annotate as objects: SalesOrderManager[Self] = SalesOrderManager()

If I try to annotate objects like so: objects: SalesOrderManager[SalesOrder] = SalesOrderManager() then I get an error that SalesOrder is not yet defined

If I try to annotate objects like so: objects: Manager[SalesOrder] = SalesOrderManager() then I get the inverse, built-in APIs returning model instances are typed correctly however the custom manager function is not found:

[rchiodo]

I'll try to answer some of your questions.

For the objects annotation, you can do this:

objects: SalesOrderManager["SalesOrder"] = SalesOrderManager()

Type annotations in a class that reference that class can put quotes around a type. It's a forward reference.

For the order object, you should be able to do this:

    order: SalesOrder = typing.cast(SalesOrder, SalesOrder.objects.first())

For the get, first, etc, those aren't yet supported. See this issue for the list of things we're thinking about adding for Django support:
#3701

[thomthom]

For the objects annotation, you can do this:

objects: SalesOrderManager["SalesOrder"] = SalesOrderManager()

I'm getting pylance errors with that:

Expected no type arguments for class "SalesOrderManager"

VSCode: 1.96.4, Pylance VSCode extension: 2024.12.1

[thomthom]

objects: models.Manager['SalesOrder'] = SalesOrderManager()

That seemed to do the trick. Is this correct?

Edit: Doing that makes QuerySets act as expected, but the custom methods on the manager is not picked up.

Edit 2:

This seems to allow Pylance to understand that the QuerySets deal with SalesOrder and pick up the manager methods:

class SalesOrderManager(models.Manager['SalesOrder']):
  def foo():
    pass

class SalesOrder
  objects: SalesOrderManager = SalesOrderManager()

[rchiodo]

Most of Django's custom methods are dynamic and Pylance doesn't know about them. It can only statically analyze your code to see methods that are allowed. Hence why the custom methods on the QuerySet are not found.

We hope to fix this eventually, but it requires something we don't currently have - a way to run python code without possibly affecting your system.

[interDist]

You should provide a type to your custom manager using a forward reference, like this:

class SalesOrderManager(models.Manager['SalesOrder'])

Then you should explicitly annotate the objects field of the model, like this:

    objects: SalesOrderManager = SalesOrderManager()

Otherwise, Pylance infers the type of the manager as BaseManager[SalesOrder] and does not find your custom methods.

[WellingtonNico]

what about the queryset results?

like this:

class MyQuerySet(models.QuerySet["MyModel"]):
    def my_filter_or_method_on_queryset(self,*args,**kwargs):
        return self.filter(...)

class MyManager(models.manager.BaseManager.from_queryset(MyQuerySet)):
    def my_filter_or_method_on_manager(self,*args,**kwargs):
        return self.filter(...)

class MyModel(models.Model):
    ...
    objects = MyManager()

this is a very basic setup and pylance can even check the types when doing filters on top of MyManager and also lints the method my_filter_or_method_on_manager, but I expect MyModel.objects.filter to return an instance of MyQuerySet, and it does that correctly, but not the correct type linting, the IDE lint for the method my_filter_or_method_on_queryset does not shows up, I need to annotate the queryset class manually every time to have access to the lint of method my_filter_or_method_on_queryset, is there a way to make BaseManager to return que queryset class type correctly?

UPDATE:
Doing like MyModel.objects.all or MyModel.objects.filter is also not returning the correct type, the point is not exactly about my manager methods, but the manager default methods to return the correct queryset type.

[rchiodo]

You could change the my_filter_or_method_on_manager to return a type. Then you wouldn't have to annotate the call sites.

[WellingtonNico]

You could change the my_filter_or_method_on_manager to return a type. Then you wouldn't have to annotate the call sites.

I missed an explanation:

doing like MyModel.objects.all or MyModel.objects.filter is also not returning the correct type, the point is not exactly about my manager methods, but the manager default methods to return the correct queryset type.

[WellingtonNico]

I endede up doing this:

[WellingtonNico]

I just solved the typing problem with this:

from typing import Type, TypeVar

from django.db import models
MODEL_TYPE = TypeVar("MODEL_TYPE", bound=models.Model)
QUERYSET_TYPE = TypeVar("QUERYSET_TYPE", bound=models.QuerySet)


class BaseDjangoManager(models.Manager, Generic[MODEL_TYPE], Generic[QUERYSET_TYPE]):

    def get_queryset(self) -> QUERYSET_TYPE:
        return super().get_queryset()

    def all(self) -> QUERYSET_TYPE:
        return super().all()

    def first(self) -> MODEL_TYPE | None:
        return super().first()

    def last(self) -> MODEL_TYPE | None:
        return super().last()

    def get(self, *args, **kwargs) -> MODEL_TYPE | None:
        return super().get(*args, **kwargs)

    def filter(self, *args, **kwargs) -> QUERYSET_TYPE:
        return super().filter(*args, **kwargs)

    def select_for_update(self, *args, **kwargs) -> QUERYSET_TYPE:
        return super().select_for_update(*args, **kwargs)

    def select_related(self, *args, **kwargs) -> QUERYSET_TYPE:
        return super().select_related(*args, **kwargs)

    def prefetch_related(self, *args, **kwargs) -> QUERYSET_TYPE:
        return super().prefetch_related(*args, **kwargs)

then I need to create my managers like this:

class ShopeeOrderManager(
    BaseDjangoManager["ShopeeOrder",ShopeeOrderQuerySet],
):
    _queryset_class = ShopeeOrderQuerySet

all lints now work just fine, no matter if manager methods or queryset methods without need to manually annotate queryset type everywhere