Check correct model on other side of many to many reverse filtering (#2283)

Co-authored-by: sobolevn <mail@sobolevn.me>

Author: flaeppe

mypy_django_plugin/lib/helpers.py

@@ -56,6 +56,7 @@ class DjangoTypeMetadata(TypedDict, total=False):
     queryset_bases: Dict[str, int]
     m2m_throughs: Dict[str, str]
     m2m_managers: Dict[str, str]
+    manager_to_model: str
 
 
 def get_django_metadata(model_info: TypeInfo) -> DjangoTypeMetadata:
@@ -94,6 +95,14 @@ def set_many_to_many_manager_info(to: TypeInfo, derived_from: str, manager_info:
     get_django_metadata(to).setdefault("m2m_managers", {})[derived_from] = manager_info.fullname
 
 
+def set_manager_to_model(manager: TypeInfo, to_model: TypeInfo) -> None:
+    get_django_metadata(manager)["manager_to_model"] = to_model.fullname
+
+
+def get_manager_to_model(manager: TypeInfo) -> Optional[str]:
+    return get_django_metadata(manager).get("manager_to_model")
+
+
 class IncompleteDefnException(Exception):
     pass
 

mypy_django_plugin/transformers/models.py

@@ -934,6 +934,7 @@ def create_many_related_manager(self, model: Instance) -> None:
         helpers.set_many_to_many_manager_info(
             to=model.type, derived_from="_default_manager", manager_info=related_manager_info
         )
+        helpers.set_manager_to_model(related_manager_info, model.type)
 
 
 class MetaclassAdjustments(ModelClassInitializer):

mypy_django_plugin/transformers/orm_lookups.py

@@ -18,7 +18,8 @@ def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext)
     if not isinstance(ctx.type, Instance) or not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
         return ctx.default_return_type
 
-    model_cls_fullname = ctx.type.args[0].type.fullname
+    manager_info = ctx.type.type
+    model_cls_fullname = helpers.get_manager_to_model(manager_info) or ctx.type.args[0].type.fullname
     model_cls = django_context.get_model_class_by_fullname(model_cls_fullname)
     if model_cls is None:
         return ctx.default_return_type

tests/typecheck/fields/test_related.yml

@@ -1480,9 +1480,9 @@
         MyModel.objects.get(xyz__isnull=False)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, others  [misc]
         MyModel.objects.exclude(xyz__isnull=False)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, others  [misc]
         other = Other()
-        other.mymodel_set.filter(xyz__isnull=True)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, mymodel, mymodel_id, other, other_id  [misc]
-        other.mymodel_set.get(xyz__isnull=True)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, mymodel, mymodel_id, other, other_id  [misc]
-        other.mymodel_set.exclude(xyz__isnull=True)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, mymodel, mymodel_id, other, other_id  [misc]
+        other.mymodel_set.filter(xyz__isnull=True)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, others  [misc]
+        other.mymodel_set.get(xyz__isnull=True)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, others  [misc]
+        other.mymodel_set.exclude(xyz__isnull=True)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, others  [misc]
         MyModel.others.through.objects.filter(xyz__isnull=False)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, mymodel, mymodel_id, other, other_id  [misc]
         MyModel.others.through.objects.get(xyz__isnull=False)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, mymodel, mymodel_id, other, other_id  [misc]
         MyModel.others.through.objects.exclude(xyz__isnull=False)  # E: Cannot resolve keyword 'xyz' into field. Choices are: id, mymodel, mymodel_id, other, other_id  [misc]
@@ -1498,3 +1498,23 @@
 
                 class MyModel(models.Model):
                     others = models.ManyToManyField(Other)
+
+-   case: test_reverse_m2m_relation_checks_other_model
+    main: |
+        from myapp.models import Author
+        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]
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+
+                class Author(models.Model):
+                    ...
+
+                class Book(models.Model):
+                    featured = models.BooleanField(default=False)
+                    authors = models.ManyToManyField(Author)