QuerySet's `.annotate()` returns an incorrect type when chained with `.values()`

[kracekumar]

Bug report

What's wrong

Iterating over the annotate result and accessing the aggregated item throws Model is not indexable.

class Author(models.Model):
    author_name = models.CharField(max_length=255)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        db_table = "authors"
        verbose_name = "Author"
        verbose_name_plural = "Authors"

    def __str__(self):
        return self.author_name

class Book(models.Model):
    BOOK_STATUS=(
        ('PUBLISHED', 'Published'),
        ('ON_HOLD', 'On Hold'),
    )
    book_name = models.CharField(max_length=255)
    author = models.ForeignKey('Author',on_delete=models.CASCADE,
                               related_name='author')
    status = models.CharField(max_length=255, choices = BOOK_STATUS,
                              default=BOOK_STATUS[0][0])
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        db_table="books"
        verbose_name="Book"
        verbose_name_plural="Books"

    def __str__(self):
        return self.book_name

class AuthorCount(TypedDict):
    authors: int

def get_book_counts_for_author() -> QuerySet[Book]:
    qs = Book.objects.values('author')
    qs1 = qs.annotate(
        authors=Count('author'))
    return qs1


def consume() -> list[int]:
    counts = []
    for item in get_book_counts_for_author():
        counts.append(item['authors']) # error happens here
    return counts

Mypy error

...
polls/models.py:79: error: Value of type "Book" is not indexable

What should be the proper way to annotate the code here other than type: ignore?

Failed Approaches

• get_book_counts_for_author() -> QuerySet[AuthorCount]:

Output:

polls/models.py:69: error: Type argument "TypedDict('polls.models.AuthorCount', {'author': builtins.int})" of "QuerySet" must be a subtype of "django.db.models.base.Model"

• get_book_counts_for_author() -> Union[QuerySet[Book], AuthorCount]:

Output:

polls/models.py:80: error: Value of type "Union[Book, str]" is not indexable
polls/models.py:80: error: Invalid index type "str" for "Union[Book, str]"; expected type "Union[int, slice]"
polls/models.py:83: error: Incompatible return value type (got "List[Union[Any, str]]", expected "List[int]")

• Tried overriding the __getitem__, __iter__ like

class BookAnnotate(QuerySet[_T]):
    def __iter__(self) -> Iterator[AuthorCount]: ...

But output was polls/models.py:65: error: Return type "Iterator[AuthorCount]" of "__iter__" incompatible with return type "Iterator[_T]" in supertype "QuerySet". The same problem happens with __getittem__. Also QuerySet[AuthorType] fails since it's not a instance of Model.

Working approach

def get_book_counts_for_author() -> Union[QuerySet[Book], AuthorCount]:
    qs = Book.objects.values('author')
    qs1 = qs.annotate(
        authors=Count('author'))
    return qs1


def consume() -> list[Optional[int]]:
    counts = []
    for item in get_book_counts_for_author():
        if isinstance(item, dict):
            counts.append(item['authors'])
    return counts

Is there a better way to annotate the code for the annotate return value. The bad part of the code is to keep checking the instance type in the consuming function and adding optional throughout the code.

How is that should be

I don't know

System information

• OS: OSX
• python version: 3.9.4
• django version: 3.2
• mypy version: 0.812
• django-stubs version: 1.8.0

[sidmitra]

@kracekumar I found some discussion around annotate on this PR.

[sobolevn]

I would love to see it rebased and merged! 👍

[kracekumar]

@sidmitra saw the PR previously and since it was not updated, felt opening a bug report to understand overall picture.

[syastrov]

I rebased the PR :)
I am not sure if it will actually fix your issue, though.

.annotate previously is typed as returning QuerySet[Any].
But the problem is that get_book_counts_for_author is returning a QuerySet where .values was called on it.
The type in django-stubs for this is called "ValuesQuerySet", but it doesn't exist in reality.
But you might be able to use return type Collection[AuthorCount] if you don't really need to use QuerySet methods on the return value.

[kracekumar]

Thank you @syastrov and a good catch on Collection. Since the return result can be used for any filtering, annotated as Union[QuerySet[Book], AuthorCount]. When it's a dictionary-like access either Collection[AuthorCount] or Iterable[Authorcount] works

[flaeppe]

For anyone interested, I just want to leave an update here that since #2319 the closest way of defining a return type is probably:

def get_book_counts_for_author() -> QuerySet[WithAnnotations[Book, AuthorCount], AuthorCount]:

That would still trigger a "Incompatible return type" error. But we need an improvement of .annotate to consider a .values having been called before we get rid of that.