diff --git a/mypy_django_plugin/transformers/querysets.py b/mypy_django_plugin/transformers/querysets.py index b65206952..1fb4e2a6a 100644 --- a/mypy_django_plugin/transformers/querysets.py +++ b/mypy_django_plugin/transformers/querysets.py @@ -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 @@ -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) @@ -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) @@ -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) @@ -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) diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index 46bcff261..25c798adb 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -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]]" + 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): @@ -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()