diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py index d49431ab1..6b4ee12aa 100644 --- 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 index c32c7235b..09167effa 100644 --- 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