Resolve queryset annotate types for expressions with static ClassVar

Instead of always typing annotated fields as Any, resolve the concrete
Python type when the expression has a static ClassVar output_field
(e.g. Count → int, Exists → bool, Length → int, Now → datetime).

Author: federicobond

mypy_django_plugin/transformers/querysets.py

@@ -16,7 +16,7 @@
 from django.db.models.sql.query import Query
 from mypy.checker import TypeChecker
 from mypy.errorcodes import NO_REDEF
-from mypy.nodes import ARG_NAMED, ARG_NAMED_OPT, ARG_STAR, CallExpr, Expression, ListExpr, SetExpr, TupleExpr
+from mypy.nodes import ARG_NAMED, ARG_NAMED_OPT, ARG_STAR, CallExpr, Expression, ListExpr, SetExpr, TupleExpr, Var
 from mypy.plugin import FunctionContext, MethodContext
 from mypy.types import AnyType, Instance, LiteralType, ProperType, TupleType, TypedDictType, TypeOfAny, get_proper_type
 from mypy.types import Type as MypyType
@@ -198,12 +198,39 @@ def gather_kwargs(ctx: MethodContext) -> dict[str, MypyType] | None:
     return kwargs
 
 
+def _resolve_output_field_type(expr_type: MypyType) -> MypyType | None:
+    """Try to resolve the Python type for an expression's output_field ClassVar.
+
+    Returns None if the output_field can't be statically resolved.
+    """
+    proper = get_proper_type(expr_type)
+    if not isinstance(proper, Instance):
+        return None
+
+    output_field_sym = proper.type.get("output_field")
+    if output_field_sym is None or output_field_sym.node is None:
+        return None
+
+    node = output_field_sym.node
+    if not isinstance(node, Var) or node.type is None:
+        return None
+
+    field_type = get_proper_type(node.type)
+    if not isinstance(field_type, Instance):
+        return None
+
+    return helpers.get_private_descriptor_type(field_type.type, "_pyi_private_get_type", is_nullable=False)
+
+
 def gather_expression_types(ctx: MethodContext) -> dict[str, MypyType]:
     kwargs = gather_kwargs(ctx)
     if not kwargs:
         return {}
 
-    # For now, we don't try to resolve the output_field of the field would be, but use Any.
+    # Try to resolve the output_field type for each expression. For expressions
+    # with a static ClassVar output_field (e.g. Count → IntegerField → int),
+    # we can infer the concrete Python type. Otherwise, fall back to Any.
+    #
     # NOTE: It's important that we use 'special_form' for 'Any' as otherwise we can
     # get stuck with mypy interpreting an overload ambiguity towards the
     # overloaded 'Field.__get__' method when its 'model' argument gets matched. This
@@ -230,7 +257,14 @@ def gather_expression_types(ctx: MethodContext) -> dict[str, MypyType]:
     # select due to the 'Any' in 'TypedDict({"foo": Any})'. But if we specify the
     # 'Any' as 'TypeOfAny.special_form' mypy doesn't consider the model instance to
     # contain 'Any' and the ambiguity goes away.
-    return {name: AnyType(TypeOfAny.special_form) for name, _ in kwargs.items()}
+    result: dict[str, MypyType] = {}
+    for name, expr_type in kwargs.items():
+        resolved = _resolve_output_field_type(expr_type)
+        if resolved is not None and not isinstance(get_proper_type(resolved), AnyType):
+            result[name] = resolved
+        else:
+            result[name] = AnyType(TypeOfAny.special_form)
+    return result
 
 
 def extract_proper_type_queryset_annotate(ctx: MethodContext, django_context: DjangoContext) -> MypyType:

tests/typecheck/managers/querysets/test_annotate.yml

@@ -561,8 +561,8 @@
 
     qs = MyModel.objects.all().annotate(num_items=Count("id"))
     obj = qs.get()
-    reveal_type(obj.num_items)  # N: Revealed type is "Any"
-    obj.nonexistent  # E: "MyModel@AnnotatedWith[TypedDict({'num_items': Any})]" has no attribute "nonexistent"  [attr-defined]
+    reveal_type(obj.num_items)  # N: Revealed type is "builtins.int"
+    obj.nonexistent  # E: "MyModel@AnnotatedWith[TypedDict({'num_items': int})]" has no attribute "nonexistent"  [attr-defined]
 
     # Custom queryset methods remain available after annotate
     qs.custom_method()
@@ -584,3 +584,49 @@
         class MyModel(models.Model):
             name = models.CharField(max_length=100)
             objects = MyManager()
+
+-   case: annotate_resolves_output_field_type
+    main: |
+        from typing_extensions import reveal_type
+        from myapp.models import User
+        from django.db.models import Count, Exists, OuterRef
+        from django.db.models.functions import Length, Now, Right, TruncDate, TruncTime
+
+        # Count has output_field: ClassVar[IntegerField] → int
+        user_count = User.objects.annotate(user_count=Count('id')).get().user_count
+        reveal_type(user_count) # N: Revealed type is "builtins.int"
+
+        # Exists has output_field: ClassVar[BooleanField] → bool
+        has_users = User.objects.annotate(has_users=Exists(User.objects.filter(pk=OuterRef('pk')))).get().has_users
+        reveal_type(has_users) # N: Revealed type is "builtins.bool"
+
+        # Length has output_field: ClassVar[IntegerField] → int
+        name_len = User.objects.annotate(name_len=Length('username')).get().name_len
+        reveal_type(name_len) # N: Revealed type is "builtins.int"
+
+        # Now has output_field: ClassVar[DateTimeField] → datetime
+        current_time = User.objects.annotate(current_time=Now()).get().current_time
+        reveal_type(current_time) # N: Revealed type is "datetime.datetime"
+
+        # Right (inherits from Left) has output_field: ClassVar[CharField] → str
+        suffix = User.objects.annotate(suffix=Right('username', 3)).get().suffix
+        reveal_type(suffix) # N: Revealed type is "builtins.str"
+
+        # TruncDate has output_field: ClassVar[DateField] → date
+        created_date = User.objects.annotate(created_date=TruncDate('created_at')).get().created_date
+        reveal_type(created_date) # N: Revealed type is "datetime.date"
+
+        # TruncTime has output_field: ClassVar[TimeField] → time
+        created_time = User.objects.annotate(created_time=TruncTime('created_at')).get().created_time
+        reveal_type(created_time) # N: Revealed type is "datetime.time"
+
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class User(models.Model):
+                    username = models.CharField(max_length=100)
+                    created_at = models.DateTimeField()