django-stubs icon indicating copy to clipboard operation
django-stubs copied to clipboard

How to correctly type manager functions that adds annotations?

Open DrakryggenLinus opened this issue 6 months ago • 0 comments

I am working with a few models where annotations are added in queryset functions. What would be the correct way to type the returned queryset to maintain the ability chain the functions while also applying typing for the annotations?

class FooQuerySet(models.QuerySet["FooModel"]):
    def with_bar(self) -> FooQuerySet:
        return self.annotate(bar="bar")
    def with_baz(self) -> FooQuerySet:
        return self.annotate(baz="baz")

FooManager = BaseManager.from_queryset(FooQuerySet)

class FooModel(models.Model):
    objects = FooManager()

Example usage: FooModel.objects.with_bar().with_baz() should return: "FooQuerySet[FooModel@AnnotatedWith[TypedDict({'bar': Any, 'baz': Any})], FooModel@AnnotatedWith[TypedDict({'bar': Any, 'baz': Any})]"

DrakryggenLinus avatar May 13 '25 10:05 DrakryggenLinus