Preserve row type when .annotate() is called on a generic QuerySet

[federicobond]

The row type was not being preserved for the TypeVar branch, so I extracted a _resolve_annotate_row_type helper to reuse on both branches.

Refs #602 (comment) (I think it may solve your issue @delfick but I haven't confirmed).

AI Policy

• I have read and agree to the AI Policy, removed any "Co-Authored-By" lines attributing coding agents, and manually reviewed the final result

[sobolevn]

Thank you!

[UnknownPlatypus]

@federicobond have you tried the test case in the issue ?

[federicobond]

Yes, I tried and they should both work now, in some cases with the right adjustments to the annotation types.

[sobolevn]

@federicobond you can contribute these tests as well, more tests in the complex areas - the better.

[delfick]

oh cool, thanks @federicobond !!!

Seems this removes 7 of the errors for me, I still have 56 left, I hopefully should be able to do some more detective work on them :)

[delfick]

Seems here is another case that shows a problem

- case: values_are_carried_through
  main: |
    from typing import Any
    from typing_extensions import reveal_type, TypeVar
    from django.db import models
    import myapp.models

    with_annotate = myapp.models.Thing.objects.all().annotate(foo=models.Max("id"))

    reveal_type(with_annotate) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Thing@AnnotatedWith[TypedDict({'foo': Any})], myapp.models.Thing@AnnotatedWith[TypedDict({'foo': Any})]]"

    with_values = with_annotate.values("id").distinct()

    reveal_type(with_values) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Thing@AnnotatedWith[TypedDict({'foo': Any})], TypedDict({'id': int})]"
  installed_apps:
    - myapp
  files:
    - path: myapp/__init__.py
    - path: myapp/models.py
      content: |
        from django.db import models
        from typing import Self

        class Thing(models.Model):
            pass

If I change with_values = with_annotate.values("id").distinct() to not have the distinct() then it passes.

edit: and seems similar for doing a slice qs[:max_sample_results]

[delfick]

also I've noticed that the .annotate will erase the custom queryset type if there is one.

[delfick]

I've run out of time to investigate, but seems for that new test case above, the fix will be something to do with

django-stubs/mypy_django_plugin/transformers/querysets.py

Line 508 in 18413d5

return default_return_type.copy_modified(args=[annotated_type, annotated_type])

[federicobond]

Yes, both cases should be fixed by #3383