Allow lookups on .annotate fields (#2376)

typeddjango · Sep 20, 2024 · eb6535f · eb6535f
1 parent f2f7bb5
commit eb6535f
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 7 deletions.
diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py
@@ -482,12 +482,12 @@ def resolve_lookup_expected_type(
         try:
             solved_lookup = self.solve_lookup_type(model_cls, lookup)
         except FieldError as exc:
-            if (
-                helpers.is_annotated_model(model_instance.type)
-                and model_instance.extra_attrs
-                and lookup in model_instance.extra_attrs.attrs
-            ):
-                return model_instance.extra_attrs.attrs[lookup]
+            if helpers.is_annotated_model(model_instance.type) and model_instance.extra_attrs:
+                # If the field comes from .annotate(), we assume Any for it
+                # and allow chaining any lookups.
+                lookup_base_field, *_ = lookup.split("__")
+                if lookup_base_field in model_instance.extra_attrs.attrs:
+                    return model_instance.extra_attrs.attrs[lookup_base_field]
 
             msg = exc.args[0]
             if model_instance.extra_attrs:

diff --git a/tests/typecheck/managers/querysets/test_annotate.yml b/tests/typecheck/managers/querysets/test_annotate.yml
@@ -257,7 +257,7 @@
                     username = models.CharField(max_length=100)
 
 
--   case: annotate_currently_allows_lookups_of_non_existent_field
+-   case: annotate_rejects_lookups_of_non_existent_field
     main: |
         from myapp.models import User
         from django.db.models.expressions import F
@@ -395,3 +395,27 @@
 
         class Entry(models.Model):
             blog = models.ForeignKey(Blog, on_delete=models.CASCADE)
+
+- case: test_annotate_allows_any_lookups_in_filter
+  main: |
+    from django.db import models
+    from myapp.models import Blog
+
+    qs = Blog.objects.annotate(distance=0)
+    reveal_type(qs)
+    reveal_type(qs.filter(distance=10))
+    reveal_type(qs.filter(distance__lt=10))
+    qs.filter(distance__lt__lt=10)
+    qs.filter(distance__unknown_lookup=10)
+  out: |
+    main:5: note: Revealed type is "django.db.models.query.QuerySet[myapp.models.Blog@AnnotatedWith[TypedDict({'distance': Any})], myapp.models.Blog@AnnotatedWith[TypedDict({'distance': Any})]]"
+    main:6: note: Revealed type is "django.db.models.query.QuerySet[myapp.models.Blog@AnnotatedWith[TypedDict({'distance': Any})], myapp.models.Blog@AnnotatedWith[TypedDict({'distance': Any})]]"
+    main:7: note: Revealed type is "django.db.models.query.QuerySet[myapp.models.Blog@AnnotatedWith[TypedDict({'distance': Any})], myapp.models.Blog@AnnotatedWith[TypedDict({'distance': Any})]]"
+  installed_apps:
+    - myapp
+  files:
+    - path: myapp/__init__.py
+    - path: myapp/models.py
+      content: |
+        from django.db import models
+        class Blog(models.Model): pass