add reverse lookups to values(), values_list()

mypy_django_plugin/django/context.py

@@ -3,7 +3,7 @@ from collections import defaultdict
 from dataclasses import dataclass
 from typing import Any, Dict, Iterator, List, Optional, TYPE_CHECKING, Tuple, Type, Sequence
 
-from django.core.exceptions import FieldError
+from django.core.exceptions import FieldError, FieldDoesNotExist
 from django.db.models.base import Model
 from django.db.models.fields.related import ForeignKey, RelatedField
 from django.utils.functional import cached_property
@@ -13,7 +13,7 @@ from pytest_mypy.utils import temp_environ
 
 from django.contrib.postgres.fields import ArrayField
 from django.db.models.fields import CharField, Field
-from django.db.models.fields.reverse_related import ForeignObjectRel
+from django.db.models.fields.reverse_related import ForeignObjectRel, ManyToOneRel, ManyToManyRel
 
 from django.db.models.sql.query import Query
 from mypy_django_plugin.lib import helpers
@@ -119,8 +119,12 @@ class DjangoLookupsContext:
                 return self.django_context.get_primary_key_field(currently_observed_model)
 
             current_field = currently_observed_model._meta.get_field(field_part)
-            if isinstance(current_field, RelatedField):
+            if isinstance(current_field, ForeignObjectRel):
                 currently_observed_model = current_field.related_model
+                current_field = self.django_context.get_primary_key_field(currently_observed_model)
+            else:
+                if isinstance(current_field, RelatedField):
+                    currently_observed_model = current_field.related_model
 
         return current_field
 

test-data/typecheck/managers/querysets/test_values.yml

@@ -88,4 +88,22 @@
                     num_articles = models.IntegerField()
                     publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE)
                 class Entry(models.Model):
-                    blog = models.ForeignKey(Blog, on_delete=models.CASCADE)
+                    blog = models.ForeignKey(Blog, on_delete=models.CASCADE)
+
+-   case: select_all_related_model_values_for_every_current_value
+    main: |
+        from myapp.models import Publisher
+        related_model_values = Publisher.objects.values('id', 'blog__name').get()
+        reveal_type(related_model_values)  # N: Revealed type is 'TypedDict({'id': builtins.int, 'blog__name': builtins.str})'
+    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):
+                    name = models.CharField(max_length=100)
+                    publisher = models.ForeignKey(Publisher, on_delete=models.CASCADE)

test-data/typecheck/managers/querysets/test_values_list.yml

@@ -28,12 +28,15 @@
 
 -   case: values_list_related_model_fields
     main: |
-        from myapp.models import Post
+        from myapp.models import Post, Blog
         values_tuple = Post.objects.values_list('blog', 'blog__num_posts', 'blog__publisher', 'blog__publisher__name').get()
         reveal_type(values_tuple[0])  # N: Revealed type is 'myapp.models.Blog'
         reveal_type(values_tuple[1])  # N: Revealed type is 'builtins.int'
         reveal_type(values_tuple[2])  # N: Revealed type is 'myapp.models.Publisher'
         reveal_type(values_tuple[3])  # N: Revealed type is 'builtins.str'
+
+        reverse_fields_list = Blog.objects.values_list('post__text').get()
+        reveal_type(reverse_fields_list)  # N: Revealed type is 'Tuple[builtins.str]'
     installed_apps:
         - myapp
     files:
@@ -47,6 +50,7 @@
                     num_posts = models.IntegerField()
                     publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE)
                 class Post(models.Model):
+                    text = models.CharField(max_length=100)
                     blog = models.ForeignKey(to=Blog, on_delete=models.CASCADE)
 
 -   case: values_list_flat_true
@@ -152,7 +156,7 @@
 
 -   case: named_true_with_related_model_fields
     main: |
-        from myapp.models import Entry
+        from myapp.models import Entry, Blog
         values = Entry.objects.values_list('blog__num_articles', 'blog__publisher__name', named=True).get()
         reveal_type(values.blog__num_articles)  # N: Revealed type is 'builtins.int'
         reveal_type(values.blog__publisher__name)  # N: Revealed type is 'builtins.str'
@@ -160,6 +164,10 @@
         pk_values = Entry.objects.values_list('blog__pk', 'blog__publisher__pk', named=True).get()
         reveal_type(pk_values.blog__pk)  # N: Revealed type is 'builtins.int'
         reveal_type(pk_values.blog__publisher__pk)  # N: Revealed type is 'builtins.int'
+
+        # reverse relation
+        reverse_values = Blog.objects.values_list('entry__text', named=True).get()
+        reveal_type(reverse_values.entry__text)  # N: Revealed type is 'builtins.str'
     installed_apps:
         - myapp
     files:
@@ -173,4 +181,5 @@
                     num_articles = models.IntegerField()
                     publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE)
                 class Entry(models.Model):
+                    text = models.CharField(max_length=100)
                     blog = models.ForeignKey(Blog, on_delete=models.CASCADE)