diff --git a/mypy_django_plugin/transformers/querysets.py b/mypy_django_plugin/transformers/querysets.py index 7145d39..3ba73b4 100644 --- a/mypy_django_plugin/transformers/querysets.py +++ b/mypy_django_plugin/transformers/querysets.py @@ -10,6 +10,7 @@ from mypy.types import AnyType, Instance from mypy.types import Type as MypyType from mypy.types import TypeOfAny +from django.db.models.fields.related import RelatedField from mypy_django_plugin.django.context import DjangoContext from mypy_django_plugin.lib import fullnames, helpers @@ -34,6 +35,9 @@ def get_field_type_from_lookup(ctx: MethodContext, django_context: DjangoContext ctx.api.fail(exc.args[0], ctx.context) return None + if isinstance(lookup_field, RelatedField) and lookup_field.column == lookup: + lookup_field = django_context.get_primary_key_field(lookup_field.related_model) + field_get_type = django_context.fields_context.get_field_get_type(helpers.get_typechecker_api(ctx), lookup_field, method=method) return field_get_type diff --git a/test-data/typecheck/managers/querysets/test_values_list.yml b/test-data/typecheck/managers/querysets/test_values_list.yml index 76d29ff..8201820 100644 --- a/test-data/typecheck/managers/querysets/test_values_list.yml +++ b/test-data/typecheck/managers/querysets/test_values_list.yml @@ -188,3 +188,20 @@ class Entry(models.Model): text = models.CharField(max_length=100) blog = models.ForeignKey(Blog, on_delete=models.CASCADE) + +- case: values_list_flat_true_with_ids + main: | + from myapp.models import Blog, Publisher + reveal_type(Blog.objects.values_list('id', flat=True)) # N: Revealed type is 'django.db.models.query.QuerySet[myapp.models.Blog, builtins.int]' + reveal_type(Blog.objects.values_list('publisher_id', flat=True)) # N: Revealed type is 'django.db.models.query.QuerySet[myapp.models.Blog, builtins.int]' + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + class Publisher(models.Model): + pass + class Blog(models.Model): + publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE) \ No newline at end of file