From 29696511066d613150089b43df827ef30a71f32f Mon Sep 17 00:00:00 2001 From: Swen Kooij Date: Mon, 24 Apr 2023 11:37:48 +0300 Subject: [PATCH] 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 --- mypy_django_plugin/lib/fullnames.py | 10 +++++- mypy_django_plugin/transformers/models.py | 4 +-- .../typecheck/models/test_related_fields.yml | 36 +++++++++++++++++++ 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/mypy_django_plugin/lib/fullnames.py b/mypy_django_plugin/lib/fullnames.py index 17027aa22..dc69d4d2d 100644 --- a/mypy_django_plugin/lib/fullnames.py +++ b/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" diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index cb7c284c2..8df69b1fe 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/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( diff --git a/tests/typecheck/models/test_related_fields.yml b/tests/typecheck/models/test_related_fields.yml index b9a801980..209838fe0 100644 --- a/tests/typecheck/models/test_related_fields.yml +++ b/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", + )