Check correct model on m2m reverse values/values_list (#2288)

Author: flaeppe

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)
 

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()