WithAnnotations does not work well with generic types (TypeVar), preventing from typing custom managers/querysets properly

[felixmeziere]

Bug report

What's wrong

First, let's establish that annotation chaining with non-custom querysets methods works:

queryset = HostingAdvert.objects.annotate(is_cool=Value(True)).annotate(is_amazing=Value(False))
reveal_type(queryset) # Revealed type is "HostingAdvertQuerySet[WithAnnotations[HostingAdvert, TypedDict({'is_cool': Any, 'is_amazing': Any})]]

The problem is that when creating a custom QuerySet containing a method that annotates it while keeping potential previous annotations (i.e. using a TypeVar for the model), the annotations are lost. Here is an example:

class SomeAnnotations(TypedDict):
    is_available: bool
    uglyness: Literal["high", "low"]

_HostingAdvertTypeVar = TypeVar("_HostingAdvertTypeVar", bound=HostingAdvert)

class GenericHostingAdvertQuerySet(QuerySet[_HostingAdvertTypeVar]):
    def annotate_with_stuff(
        self: GenericHostingAdvertQuerySet[_HostingAdvertTypeVar],
    ) -> GenericHostingAdvertQuerySet[WithAnnotations[_HostingAdvertTypeVar, SomeAnnotations]]:
        return self.annotate(is_available=Exists(...), ugliness=Value("high"))

class HostingAdvertQuerySet(GenericHostingAdvertQuerySet[HostingAdvert]):
    pass

queryset = HostingAdvertQuerySet().annotate(is_cool=Value(True)).annotate_with_stuff()
reveal_type(queryset) # INCORRECT: Revealed type is "GenericHostingAdvertQuerySet[HostingAdvert]" . All annotations have been lost.

If instead of using a TypeVar I use the concrete HostingAdvert type, then the annotations of the method work but any previous annotations on the queryset will be lost which is not satisfactory:

class SomeAnnotations(TypedDict):
    is_available: bool
    uglyness: Literal["high", "low"]

_HostingAdvertTypeVar = TypeVar("_HostingAdvertTypeVar", bound=HostingAdvert)

class HostingAdvertQuerySet(QuerySet[_HostingAdvertTypeVar]):
    def annotate_with_stuff(self) -> HostingAdvertQuerySet[WithAnnotations[HostingAdvert, SomeAnnotations]]:
        return self.annotate(is_available=Exists(...), ugliness=Value("high"))

class HostingAdvertQuerySet(GenericHostingAdvertQuerySet[HostingAdvert]):
    pass

queryset = HostingAdvertQuerySet().annotate(is_cool=Value(True)).annotate_with_stuff()
reveal_type(queryset) # NOT WHAT WE WANT: Revealed type is GenericHostingAdvertQuerySet[WithAnnotations[HostingAdvert, TypedDict('SomeAnnotations', {'is_available': builtins.bool, 'uglyness': Union[Literal['high'], Literal['low']]})]] meaning the is_cool annotation has been lost.

How should it be instead

Revealed type should be

GenericHostingAdvertQuerySet[WithAnnotations[HostingAdvert, TypedDict('SomeAnnotations', {'is_cool': Any, 'is_available': builtins.bool, 'uglyness': Union[Literal['high'], Literal['low']]})]]

which keeps the previous annotations and adds new ones.

System information

• OS: Mac OS Monterey
• python version: 10.0.3
• django version: 4.0.4
• mypy version: 0.961
• django-stubs version: current master (so just after 1.12.0)
• django-stubs-ext version: 0.5.0

[felixmeziere]

@syastrov any ideas on this one too?

[pablo-tx]

@felixmeziere did you manage to find a solution?

[felixmeziere]

I think I found a workaround using Self in the return type of some functions