Preserve row type in merge_annotations_from_custom_method (#3383)

Author: federicobond

mypy_django_plugin/transformers/querysets.py

@@ -505,7 +505,16 @@ def merge_annotations_from_custom_method(ctx: MethodContext, django_context: Dja
 
     api = helpers.get_typechecker_api(ctx)
     annotated_type = get_annotated_type(api, django_model.typ, fields_dict=new_td)
-    return default_return_type.copy_modified(args=[annotated_type, annotated_type])
+    new_args: list[MypyType] = [annotated_type]
+    if len(default_return_type.args) > 1:
+        original_row = get_proper_type(default_return_type.args[1])
+        if isinstance(original_row, Instance) and helpers.is_model_type(original_row.type):
+            new_args.append(annotated_type)
+        else:
+            new_args.append(default_return_type.args[1])
+    else:
+        new_args.append(annotated_type)
+    return default_return_type.copy_modified(args=new_args)
 
 
 def resolve_field_lookups(lookup_exprs: Sequence[Expression], django_context: DjangoContext) -> list[str] | None:

tests/typecheck/managers/querysets/test_annotate.yml

@@ -362,6 +362,33 @@
                     num_posts = models.IntegerField()
                     text = models.CharField(max_length=100)
 
+-   case: annotate_then_values_list_flat_slice_preserves_row_type
+    main: |
+        from typing import Iterable
+        from typing_extensions import reveal_type
+        from myapp.models import Blog
+        from django.db.models import F
+
+        qs = Blog.objects.annotate(a=F("id")).values_list("id", flat=True)
+        reveal_type(qs)  # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Blog@AnnotatedWith[TypedDict({'a': Any})], int]"
+        reveal_type(qs[:])  # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Blog@AnnotatedWith[TypedDict({'a': Any})], int]"
+
+        def get_blog_ids() -> Iterable[int]:
+            return Blog.objects.annotate(a=F("id")).values_list("id", flat=True)[:]
+
+        # Also preserve non-flat tuple row types and TypedDict row types from values().
+        reveal_type(Blog.objects.annotate(a=F("id")).values_list("id", "text")[:])  # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Blog@AnnotatedWith[TypedDict({'a': Any})], tuple[int, str]]"
+        reveal_type(Blog.objects.annotate(a=F("id")).values("id")[:])  # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Blog@AnnotatedWith[TypedDict({'a': Any})], TypedDict({'id': int})]"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class Blog(models.Model):
+                    text = models.CharField(max_length=100)
+
 - case: test_annotate_with_filter
   main: |
     from django.db import models