values_list for related model id flat True (#113)

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

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)