Type defined with WithAnnotations fails mypy 0.981

[AllexVeldman]

Bug report

What's wrong

Types defined using WithAnnotations started failing with mypy 0.981.

from typing import TypedDict
from django.db.models import IntegerField, Model, QuerySet, Value
from django_stubs_ext import WithAnnotations

class Bar(Model):
    x = IntegerField()

class BarAnnotations(TypedDict):
    y: int

def foo(xyz: QuerySet[Bar]) -> QuerySet[WithAnnotations[Bar, BarAnnotations]]:
    return xyz.annotate(y=Value(2))

reveal_type(foo)

mypy 0.971:

tmp.py:19: note: Revealed type is "def (xyz: django.db.models.query._QuerySet[tmp.Bar, tmp.Bar]) -> django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[tmp__Bar, TypedDict('tmp.BarAnnotations', {'y': builtins.int})], django_stubs_ext.WithAnnotations[tmp__Bar, TypedDict('tmp.BarAnnotations', {'y': builtins.int})]]"
Success: no issues found in 1 source file

mypy 0.981/0.982:

tmp.py:16: error: Incompatible return value type (got "_QuerySet[WithAnnotations[tmp__Bar, TypedDict({'y': Any})], WithAnnotations[tmp__Bar, TypedDict({'y': Any})]]", expected "_QuerySet[WithAnnotations[tmp__Bar], WithAnnotations[tmp__Bar]]")  [return-value]
tmp.py:19: note: Revealed type is "def (xyz: django.db.models.query._QuerySet[tmp.Bar, tmp.Bar]) -> django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[tmp__Bar], django_stubs_ext.WithAnnotations[tmp__Bar]]"
Found 1 error in 1 file (checked 1 source file)

It seems mypy infers the return-type of foo as un-annotated, I tried to dig into how the WithAnnotations works but got lost in the plugin quite quickly. As I'm not sure if this is a mypy or a django-stubs issue, some pointers on where to dig would be welcome.

System information

• OS: macOS 12.6
• python version: 3.10.2
• django version: 4.1.2
• mypy version: 0.982
• django-stubs version: 1.12.0
• django-stubs-ext version: 0.5.0

[AllexVeldman]

from typing import TypedDict

from django.db.models import IntegerField, Model, QuerySet, Value
from django_stubs_ext import WithAnnotations


class Bar(Model):
    x = IntegerField()

class BarAnnotations(TypedDict):
    y: int

def with_param(x: QuerySet[Bar]) -> QuerySet[WithAnnotations[Bar, BarAnnotations]]:
    return x.annotate(y=Value(2))
    # error: Incompatible return value type (got "_QuerySet[WithAnnotations[tmp__Bar, TypedDict({'y': Any})], WithAnnotations[tmp__Bar, TypedDict({'y': Any})]]", expected "_QuerySet[WithAnnotations[tmp__Bar], WithAnnotations[tmp__Bar]]")  [return-value]

def without_param() -> QuerySet[WithAnnotations[Bar, BarAnnotations]]:
    qs = Bar.objects.all()
    return qs.annotate(y=Value(2))
    # OK

def without_param_with_type() -> QuerySet[WithAnnotations[Bar, BarAnnotations]]:
    qs: QuerySet[Bar] = Bar.objects.all()
    return qs.annotate(y=Value(2))
    # error: Incompatible return value type (got "_QuerySet[WithAnnotations[tmp__Bar, TypedDict({'y': Any})], WithAnnotations[tmp__Bar, TypedDict({'y': Any})]]", expected "_QuerySet[WithAnnotations[tmp__Bar], WithAnnotations[tmp__Bar]]")  [return-value]

reveal_type(with_param)
tmp.py:32: note: Revealed type is "def (x: django.db.models.query._QuerySet[Any, Any]) -> django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[tmp__Bar], django_stubs_ext.WithAnnotations[tmp__Bar]]"
reveal_type(without_param)
tmp.py:33: note: Revealed type is "def () -> django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[tmp__Bar], django_stubs_ext.WithAnnotations[tmp__Bar]]"
reveal_type(without_param_with_type)
tmp.py:34: note: Revealed type is "def () -> django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[tmp__Bar], django_stubs_ext.WithAnnotations[tmp__Bar]]"

[MrkGrgsn]

FYI, I couldn't replicate this with the following setup:

• OS: Fedora 38
• python version: 3.11.6
• django version: 3.2.23
• mypy version: 1.7.1
• django-stubs version: 4.2.7
• django-stubs-ext version: 4.2.7

mypy complained about Bar.objects not being defined and I tried with objects = Manager() but no other warnings were raised in either case.

tmp.py:32: note: Revealed type is "def (x: django.db.models.query._QuerySet[mancini.type_example.tmp.Bar, mancini.type_example.tmp.Bar]) -> django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[tmp__Bar, TypedDict('mancini.type_example.tmp.BarAnnotations', {'y': builtins.int})], django_stubs_ext.WithAnnotations[tmp__Bar, TypedDict('mancini.type_example.tmp.BarAnnotations', {'y': builtins.int})]]"
tmp.py:33: note: Revealed type is "def () -> django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[tmp__Bar, TypedDict('mancini.type_example.tmp.BarAnnotations', {'y': builtins.int})], django_stubs_ext.WithAnnotations[tmp__Bar, TypedDict('mancini.type_example.tmp.BarAnnotations', {'y': builtins.int})]]"
tmp.py:34: note: Revealed type is "def () -> django.db.models.query._QuerySet[django_stubs_ext.WithAnnotations[tmp__Bar, TypedDict('mancini.type_example.tmp.BarAnnotations', {'y': builtins.int})], django_stubs_ext.WithAnnotations[tmp__Bar, TypedDict('mancini.type_example.tmp.BarAnnotations', {'y': builtins.int})]]"

[MrkGrgsn]

I noted the flaw in my testing of this (not setting objects to the custom manager), so reworked the OP's example slightly (and adding a couple of variants to test) as follows:

from __future__ import annotations

from typing import TYPE_CHECKING, TypedDict

from django.db.models import IntegerField, Model, Value, QuerySet
from django_stubs_ext import WithAnnotations

class BarAnnotations(TypedDict):
    y: int

class BarQuerySet(QuerySet["Bar"]):

    def custom(self) -> QuerySet[WithAnnotations["Bar", BarAnnotations]]:
        return self.annotate(y=1)
       # error: Incompatible return value type (got "BarQuerySet[WithAnnotations[mancini__type_example__models__Bar, TypedDict({'y': Any})], WithAnnotations[mancini__type_example__models__Bar, TypedDict({'y': Any})]]", expected "_QuerySet[WithAnnotations[mancini__type_example__models__Bar, TypedDict('mancini.type_example.models.BarAnnotations', {'y': builtins.int})], WithAnnotations[mancini__type_example__models__Bar, TypedDict('mancini.type_example.models.BarAnnotations', {'y': builtins.int})]]")  [return-value]

def with_param(x: QuerySet["Bar"]) -> QuerySet[WithAnnotations["Bar", BarAnnotations]]:
    return x.annotate(y=Value(2))
    # OK

def with_param_custom(x: BarQuerySet) -> QuerySet[WithAnnotations["Bar", BarAnnotations]]:
    return x.annotate(y=Value(2))
    # error: Incompatible return value type (got "BarQuerySet[WithAnnotations[mancini__type_example__models__Bar, TypedDict({'y': Any})], WithAnnotations[mancini__type_example__models__Bar, TypedDict({'y': Any})]]", expected "_QuerySet[WithAnnotations[mancini__type_example__models__Bar, TypedDict('mancini.type_example.models.BarAnnotations', {'y': builtins.int})], WithAnnotations[mancini__type_example__models__Bar, TypedDict('mancini.type_example.models.BarAnnotations', {'y': builtins.int})]]")  [return-value]

def without_param() -> QuerySet[WithAnnotations["Bar", BarAnnotations]]:
    qs = Bar.objects.all()
    return qs.annotate(y=Value(2))
    # error: Incompatible return value type (got "BarQuerySet[WithAnnotations[mancini__type_example__models__Bar, TypedDict({'y': Any})], WithAnnotations[mancini__type_example__models__Bar, TypedDict({'y': Any})]]", expected "_QuerySet[WithAnnotations[mancini__type_example__models__Bar, TypedDict('mancini.type_example.models.BarAnnotations', {'y': builtins.int})], WithAnnotations[mancini__type_example__models__Bar, TypedDict('mancini.type_example.models.BarAnnotations', {'y': builtins.int})]]")  [return-value]

def without_param_with_type() -> QuerySet[WithAnnotations["Bar", BarAnnotations]]:
    qs: QuerySet["Bar"] = Bar.objects.all()
    return qs.annotate(y=Value(2))
    # OK

class Bar(Model):
    x = IntegerField()

    objects = BarQuerySet.as_manager()

In my environment, mypy complains only when the queryset to be annotated is typed with the custom queryset class, presumably because BarQuerySet doesn't inherit from _QuerySet.

NB. It doesn't seem totally correct that mypy reports a parameterized form of BarQuerySet given it's not a generic type. I'm unsure if that contributes to this specific problem though: https://github.com/typeddjango/django-stubs/pull/1030#issuecomment-1183628183