Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Check correct model on m2m reverse values/values_list #2288

Merged
merged 2 commits into from
Jul 29, 2024
Merged

Check correct model on m2m reverse values/values_list #2288

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d824cde
Check correct model on m2m reverse values/values_list
flaeppe Jul 29, 2024
cb7d2fa
fixup! Check correct model on m2m reverse values/values_list
flaeppe Jul 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 15 additions & 4 deletions mypy_django_plugin/transformers/querysets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,12 @@
from django.db.models.base import Model
from django.db.models.fields.related import RelatedField
from django.db.models.fields.reverse_related import ForeignObjectRel
from mypy.checker import TypeChecker
from mypy.nodes import ARG_NAMED, ARG_NAMED_OPT, Expression
from mypy.plugin import FunctionContext, MethodContext
from mypy.types import AnyType, Instance, TupleType, TypedDictType, TypeOfAny, get_proper_type
from mypy.types import Type as MypyType
from mypy.typevars import fill_typevars

from mypy_django_plugin.django.context import DjangoContext, LookupsAreUnsupported
from mypy_django_plugin.lib import fullnames, helpers
Expand All @@ -17,7 +19,16 @@
from mypy_django_plugin.transformers.models import get_or_create_annotated_type


def _extract_model_type_from_queryset(queryset_type: Instance) -> Optional[Instance]:
def _extract_model_type_from_queryset(queryset_type: Instance, api: TypeChecker) -> Optional[Instance]:
if queryset_type.type.has_base(fullnames.MANAGER_CLASS_FULLNAME):
to_model_fullname = helpers.get_manager_to_model(queryset_type.type)
if to_model_fullname is not None:
to_model = helpers.lookup_fully_qualified_typeinfo(api, to_model_fullname)
if to_model is not None:
to_model_instance = fill_typevars(to_model)
assert isinstance(to_model_instance, Instance)
return to_model_instance

for base_type in [queryset_type, *queryset_type.type.bases]:
if (
len(base_type.args)
Expand Down Expand Up @@ -161,7 +172,7 @@ def extract_proper_type_queryset_values_list(ctx: MethodContext, django_context:
if not isinstance(default_return_type, Instance):
return ctx.default_return_type

model_type = _extract_model_type_from_queryset(ctx.type)
model_type = _extract_model_type_from_queryset(ctx.type, helpers.get_typechecker_api(ctx))
if model_type is None:
return AnyType(TypeOfAny.from_omitted_generics)

Expand Down Expand Up @@ -221,7 +232,7 @@ def extract_proper_type_queryset_annotate(ctx: MethodContext, django_context: Dj
if not isinstance(default_return_type, Instance):
return ctx.default_return_type

model_type = _extract_model_type_from_queryset(ctx.type)
model_type = _extract_model_type_from_queryset(ctx.type, helpers.get_typechecker_api(ctx))
if model_type is None:
return AnyType(TypeOfAny.from_omitted_generics)

Expand Down Expand Up @@ -288,7 +299,7 @@ def extract_proper_type_queryset_values(ctx: MethodContext, django_context: Djan
if not isinstance(default_return_type, Instance):
return ctx.default_return_type

model_type = _extract_model_type_from_queryset(ctx.type)
model_type = _extract_model_type_from_queryset(ctx.type, helpers.get_typechecker_api(ctx))
if model_type is None:
return AnyType(TypeOfAny.from_omitted_generics)

Expand Down
31 changes: 31 additions & 0 deletions tests/typecheck/fields/test_related.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1502,14 +1502,36 @@
- case: test_reverse_m2m_relation_checks_other_model
main: |
from myapp.models import Author
# With builtin manager/queryset
Author().book_set.filter(featured=True)
Author().book_set.filter(xyz=True) # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]
reveal_type(Author().book_set.values_list("featured")) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Book, Tuple[builtins.bool]]"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rare case, where Tuple[bool] and not Tuple[bool, ...] is correct.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah we're fetching a single column so we get a tuple with a single bool

Author().book_set.values_list("xyz") # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]
reveal_type(Author().book_set.values("featured")) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Book, TypedDict({'featured': builtins.bool})]"
Author().book_set.values("xyz") # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]
reveal_type(Author().book_set.filter(featured=True).values_list("featured")) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Book, Tuple[builtins.bool]]"
Author().book_set.filter(featured=True).values_list("xyz") # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]
reveal_type(Author().book_set.filter(featured=True).values("featured")) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Book, TypedDict({'featured': builtins.bool})]"
Author().book_set.filter(featured=True).values("xyz") # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]

# With a custom manager/queryset
Author().other_set.filter(featured=True)
Author().other_set.filter(xyz=True) # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]
reveal_type(Author().other_set.values_list("featured")) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Other, Tuple[builtins.bool]]"
Author().other_set.values_list("xyz") # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]
reveal_type(Author().other_set.values("featured")) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Other, TypedDict({'featured': builtins.bool})]"
Author().other_set.values("xyz") # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]
reveal_type(Author().other_set.filter(featured=True).values_list("featured")) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Other, Tuple[builtins.bool]]"
Author().other_set.filter(featured=True).values_list("xyz") # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]
reveal_type(Author().other_set.filter(featured=True).values("featured")) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Other, TypedDict({'featured': builtins.bool})]"
Author().other_set.filter(featured=True).values("xyz") # E: Cannot resolve keyword 'xyz' into field. Choices are: authors, featured, id [misc]
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from typing_extensions import Self
from django.db import models

class Author(models.Model):
Expand All @@ -1518,3 +1540,12 @@
class Book(models.Model):
featured = models.BooleanField(default=False)
authors = models.ManyToManyField(Author)

class OtherQuerySet(models.QuerySet["Other"]):
def custom(self) -> Self: ...

OtherManager = models.Manager.from_queryset(OtherQuerySet)
class Other(models.Model):
featured = models.BooleanField()
authors = models.ManyToManyField(Author)
objects = OtherManager()
Loading