From 4f935edd475e56cebdf0c7edff4d3000f5ec2849 Mon Sep 17 00:00:00 2001 From: Maxim Kurnikov Date: Thu, 18 Jul 2019 03:47:50 +0300 Subject: [PATCH] add reverse lookups to values(), values_list() --- mypy_django_plugin/django/context.py | 10 +++++++--- .../managers/querysets/test_values.yml | 20 ++++++++++++++++++- .../managers/querysets/test_values_list.yml | 13 ++++++++++-- 3 files changed, 37 insertions(+), 6 deletions(-) diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py index 9d4b14d..2d73888 100644 --- a/mypy_django_plugin/django/context.py +++ b/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 diff --git a/test-data/typecheck/managers/querysets/test_values.yml b/test-data/typecheck/managers/querysets/test_values.yml index 1632abc..420fcd0 100644 --- a/test-data/typecheck/managers/querysets/test_values.yml +++ b/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) \ No newline at end of file + 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) diff --git a/test-data/typecheck/managers/querysets/test_values_list.yml b/test-data/typecheck/managers/querysets/test_values_list.yml index 1196547..26cc3da 100644 --- a/test-data/typecheck/managers/querysets/test_values_list.yml +++ b/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)