Allow any reverse relation using ForeignObjectRel to be type checked (#1451)

* Allow any reverse relation using ForeignObjectRel to be type checked

Currently, reverse relations created using just `ForeignObjectRel`
aren't checked. A real-world example of this is `django-composite-foreignkey`.
Reverse relations created by it aren't discovered by the mypy plugin.

* Add test demonstrating support for multi-column ForeignObject

* Allow foreign keys using `ForeignObject` to be type checked

* Freeze `RELATED_FIELDS_CLASSES` to prevent accidental mutation

Author: Photonios

mypy_django_plugin/lib/fullnames.py

@@ -6,6 +6,7 @@
 ARRAY_FIELD_FULLNAME = "django.contrib.postgres.fields.array.ArrayField"
 AUTO_FIELD_FULLNAME = "django.db.models.fields.AutoField"
 GENERIC_FOREIGN_KEY_FULLNAME = "django.contrib.contenttypes.fields.GenericForeignKey"
+FOREIGN_OBJECT_FULLNAME = "django.db.models.fields.related.ForeignObject"
 FOREIGN_KEY_FULLNAME = "django.db.models.fields.related.ForeignKey"
 ONETOONE_FIELD_FULLNAME = "django.db.models.fields.related.OneToOneField"
 MANYTOMANY_FIELD_FULLNAME = "django.db.models.fields.related.ManyToManyField"
@@ -30,7 +31,14 @@
     BASE_MANAGER_CLASS_FULLNAME,
 }
 
-RELATED_FIELDS_CLASSES = {FOREIGN_KEY_FULLNAME, ONETOONE_FIELD_FULLNAME, MANYTOMANY_FIELD_FULLNAME}
+RELATED_FIELDS_CLASSES = frozenset(
+    (
+        FOREIGN_OBJECT_FULLNAME,
+        FOREIGN_KEY_FULLNAME,
+        ONETOONE_FIELD_FULLNAME,
+        MANYTOMANY_FIELD_FULLNAME,
+    )
+)
 
 MIGRATION_CLASS_FULLNAME = "django.db.migrations.migration.Migration"
 OPTIONS_CLASS_FULLNAME = "django.db.models.options.Options"

mypy_django_plugin/transformers/models.py

@@ -3,7 +3,7 @@
 from django.db.models import Manager, Model
 from django.db.models.fields import DateField, DateTimeField, Field
 from django.db.models.fields.related import ForeignKey
-from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel, OneToOneRel
+from django.db.models.fields.reverse_related import ForeignObjectRel, OneToOneRel
 from mypy.checker import TypeChecker
 from mypy.nodes import ARG_STAR2, Argument, AssignmentStmt, CallExpr, Context, NameExpr, TypeInfo, Var
 from mypy.plugin import AnalyzeTypeContext, AttributeContext, CheckerPluginInterface, ClassDefContext
@@ -464,7 +464,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                 self.add_new_node_to_model_class(attname, Instance(related_model_info, []))
                 continue
 
-            if isinstance(relation, (ManyToOneRel, ManyToManyRel)):
+            if isinstance(relation, ForeignObjectRel):
                 related_manager_info = None
                 try:
                     related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(

tests/typecheck/models/test_related_fields.yml

@@ -87,3 +87,39 @@
               b = models.ForeignKey(Model2, related_name="test2", on_delete=models.CASCADE)
 
               objects = Model4Manager()
+
+-   case: test_related_name_foreign_object_multi_column
+    main: |
+      from app1.models import Model1, Model2
+
+      reveal_type(Model2.model_1) # N: Revealed type is "django.db.models.fields.related.ForeignObject[app1.models.Model1, app1.models.Model1]"
+      reveal_type(Model2().model_1) # N: Revealed type is "app1.models.Model1"
+      reveal_type(Model1.model_2s) # N: Revealed type is "django.db.models.manager.RelatedManager[app1.models.Model2]"
+      reveal_type(Model1().model_2s) # N: Revealed type is "django.db.models.manager.RelatedManager[app1.models.Model2]"
+
+    installed_apps:
+      - app1
+    files:
+      - path: app1/__init__.py
+      - path: app1/models.py
+        content: |
+          from django.db import models
+          from django.db.models.fields.related import ForeignObject
+
+          class Model1(models.Model):
+              type = models.TextField()
+              ref = models.TextField()
+
+          class Model2(models.Model):
+              name = models.TextField()
+
+              model_1_type = models.TextField()
+              model_2_ref = models.TextField()
+
+              model_1 = ForeignObject(
+                  Model1,
+                  to_fields=["type", "ref"],
+                  from_fields=["model_1_type", "model_2_ref"],
+                  on_delete=models.CASCADE,
+                  related_name="model_2s",
+              )