Skip to content

Commit 31e7950

Browse files
authored
values(), values_list() with ManyToManyField (#267)
1 parent 0cba3f9 commit 31e7950

File tree

4 files changed

+50
-12
lines changed

4 files changed

+50
-12
lines changed

mypy_django_plugin/django/context.py

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ def get_model_relations(self, model_cls: Type[Model]) -> Iterator[ForeignObjectR
124124
if isinstance(field, ForeignObjectRel):
125125
yield field
126126

127-
def get_field_lookup_exact_type(self, api: TypeChecker, field: Field) -> MypyType:
127+
def get_field_lookup_exact_type(self, api: TypeChecker, field: Union[Field, ForeignObjectRel]) -> MypyType:
128128
if isinstance(field, (RelatedField, ForeignObjectRel)):
129129
related_model_cls = field.related_model
130130
primary_key_field = self.get_primary_key_field(related_model_cls)
@@ -134,10 +134,8 @@ def get_field_lookup_exact_type(self, api: TypeChecker, field: Field) -> MypyTyp
134134
if rel_model_info is None:
135135
return AnyType(TypeOfAny.explicit)
136136

137-
model_and_primary_key_type = UnionType.make_union([Instance(rel_model_info, []),
138-
primary_key_type])
137+
model_and_primary_key_type = UnionType.make_union([Instance(rel_model_info, []), primary_key_type])
139138
return helpers.make_optional(model_and_primary_key_type)
140-
# return helpers.make_optional(Instance(rel_model_info, []))
141139

142140
field_info = helpers.lookup_class_typeinfo(api, field.__class__)
143141
if field_info is None:
@@ -228,21 +226,22 @@ def get_attname(self, field: Field) -> str:
228226
attname = field.attname
229227
return attname
230228

231-
def get_field_nullability(self, field: Field, method: Optional[str]) -> bool:
229+
def get_field_nullability(self, field: Union[Field, ForeignObjectRel], method: Optional[str]) -> bool:
232230
nullable = field.null
233231
if not nullable and isinstance(field, CharField) and field.blank:
234232
return True
235233
if method == '__init__':
236-
if field.primary_key or isinstance(field, ForeignKey):
234+
if ((isinstance(field, Field) and field.primary_key)
235+
or isinstance(field, ForeignKey)):
237236
return True
238237
if method == 'create':
239238
if isinstance(field, AutoField):
240239
return True
241-
if field.has_default():
240+
if isinstance(field, Field) and field.has_default():
242241
return True
243242
return nullable
244243

245-
def get_field_set_type(self, api: TypeChecker, field: Field, *, method: str) -> MypyType:
244+
def get_field_set_type(self, api: TypeChecker, field: Union[Field, ForeignObjectRel], *, method: str) -> MypyType:
246245
""" Get a type of __set__ for this specific Django field. """
247246
target_field = field
248247
if isinstance(field, ForeignKey):
@@ -259,7 +258,7 @@ def get_field_set_type(self, api: TypeChecker, field: Field, *, method: str) ->
259258
field_set_type = helpers.convert_any_to_type(field_set_type, argument_field_type)
260259
return field_set_type
261260

262-
def get_field_get_type(self, api: TypeChecker, field: Field, *, method: str) -> MypyType:
261+
def get_field_get_type(self, api: TypeChecker, field: Union[Field, ForeignObjectRel], *, method: str) -> MypyType:
263262
""" Get a type of __get__ for this specific Django field. """
264263
field_info = helpers.lookup_class_typeinfo(api, field.__class__)
265264
if field_info is None:
@@ -303,7 +302,10 @@ def get_field_related_model_cls(self, field: Union[RelatedField, ForeignObjectRe
303302

304303
return related_model_cls
305304

306-
def _resolve_field_from_parts(self, field_parts: Iterable[str], model_cls: Type[Model]) -> Field:
305+
def _resolve_field_from_parts(self,
306+
field_parts: Iterable[str],
307+
model_cls: Type[Model]
308+
) -> Union[Field, ForeignObjectRel]:
307309
currently_observed_model = model_cls
308310
field = None
309311
for field_part in field_parts:
@@ -325,7 +327,7 @@ def _resolve_field_from_parts(self, field_parts: Iterable[str], model_cls: Type[
325327
assert field is not None
326328
return field
327329

328-
def resolve_lookup_into_field(self, model_cls: Type[Model], lookup: str) -> Field:
330+
def resolve_lookup_into_field(self, model_cls: Type[Model], lookup: str) -> Union[Field, ForeignObjectRel]:
329331
query = Query(model_cls)
330332
lookup_parts, field_parts, is_expression = query.solve_lookup_type(lookup)
331333
if lookup_parts:

mypy_django_plugin/transformers/querysets.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from django.core.exceptions import FieldError
55
from django.db.models.base import Model
66
from django.db.models.fields.related import RelatedField
7+
from django.db.models.fields.reverse_related import ForeignObjectRel
78
from mypy.nodes import Expression, NameExpr
89
from mypy.plugin import FunctionContext, MethodContext
910
from mypy.types import AnyType, Instance
@@ -47,7 +48,8 @@ def get_field_type_from_lookup(ctx: MethodContext, django_context: DjangoContext
4748
except LookupsAreUnsupported:
4849
return AnyType(TypeOfAny.explicit)
4950

50-
if isinstance(lookup_field, RelatedField) and lookup_field.column == lookup:
51+
if ((isinstance(lookup_field, RelatedField) and lookup_field.column == lookup)
52+
or isinstance(lookup_field, ForeignObjectRel)):
5153
related_model_cls = django_context.get_field_related_model_cls(lookup_field)
5254
if related_model_cls is None:
5355
return AnyType(TypeOfAny.from_error)

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,3 +107,20 @@
107107
class Blog(models.Model):
108108
name = models.CharField(max_length=100)
109109
publisher = models.ForeignKey(Publisher, on_delete=models.CASCADE)
110+
111+
- case: values_of_many_to_many_field
112+
main: |
113+
from myapp.models import Author, Book
114+
reveal_type(Book.objects.values('authors')) # N: Revealed type is 'django.db.models.query.ValuesQuerySet[myapp.models.Book, TypedDict({'authors': builtins.int})]'
115+
reveal_type(Author.objects.values('books')) # N: Revealed type is 'django.db.models.query.ValuesQuerySet[myapp.models.Author, TypedDict({'books': builtins.int})]'
116+
installed_apps:
117+
- myapp
118+
files:
119+
- path: myapp/__init__.py
120+
- path: myapp/models.py
121+
content: |
122+
from django.db import models
123+
class Author(models.Model):
124+
pass
125+
class Book(models.Model):
126+
authors = models.ManyToManyField(Author, related_name='books')

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,3 +224,20 @@
224224
pass
225225
class Transaction(models.Model):
226226
total = models.IntegerField()
227+
228+
- case: values_list_of_many_to_many_field
229+
main: |
230+
from myapp.models import Author, Book
231+
reveal_type(Book.objects.values_list('authors')) # N: Revealed type is 'django.db.models.query.ValuesQuerySet[myapp.models.Book, Tuple[builtins.int]]'
232+
reveal_type(Author.objects.values_list('books')) # N: Revealed type is 'django.db.models.query.ValuesQuerySet[myapp.models.Author, Tuple[builtins.int]]'
233+
installed_apps:
234+
- myapp
235+
files:
236+
- path: myapp/__init__.py
237+
- path: myapp/models.py
238+
content: |
239+
from django.db import models
240+
class Author(models.Model):
241+
pass
242+
class Book(models.Model):
243+
authors = models.ManyToManyField(Author, related_name='books')

0 commit comments

Comments
 (0)