django-stubs
django-stubs copied to clipboard
How to correctly type manager functions that adds annotations?
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})]"