Add get_models_foreign_key helper

Author: UnknownPlatypus

mypy_django_plugin/django/context.py

@@ -118,6 +118,11 @@ def get_model_fields(self, model_cls: Type[Model]) -> Iterator["Field[Any, Any]"
             if isinstance(field, Field):
                 yield field
 
+    def get_model_foreign_keys(self, model_cls: Type[Model]) -> Iterator["ForeignKey[Any, Any]"]:
+        for field in model_cls._meta.get_fields():
+            if isinstance(field, ForeignKey):
+                yield field
+
     def get_model_relations(self, model_cls: Type[Model]) -> Iterator[ForeignObjectRel]:
         for field in model_cls._meta.get_fields():
             if isinstance(field, ForeignObjectRel):

mypy_django_plugin/transformers/models.py

@@ -2,7 +2,6 @@
 
 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 ForeignObjectRel, OneToOneRel
 from mypy.checker import TypeChecker
 from mypy.nodes import ARG_STAR2, Argument, AssignmentStmt, CallExpr, Context, NameExpr, TypeInfo, Var
@@ -234,41 +233,40 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
 
 class AddRelatedModelsId(ModelClassInitializer):
     def run_with_model_cls(self, model_cls: Type[Model]) -> None:
-        for field in model_cls._meta.get_fields():
-            if isinstance(field, ForeignKey):
-                related_model_cls = self.django_context.get_field_related_model_cls(field)
-                if related_model_cls is None:
-                    error_context: Context = self.ctx.cls
-                    field_sym = self.ctx.cls.info.get(field.name)
-                    if field_sym is not None and field_sym.node is not None:
-                        error_context = field_sym.node
-                    self.api.fail(
-                        f"Cannot find model {field.related_model!r} referenced in field {field.name!r}",
-                        ctx=error_context,
-                    )
-                    self.add_new_node_to_model_class(field.attname, AnyType(TypeOfAny.explicit))
-                    continue
+        for field in self.django_context.get_model_foreign_keys(model_cls):
+            related_model_cls = self.django_context.get_field_related_model_cls(field)
+            if related_model_cls is None:
+                error_context: Context = self.ctx.cls
+                field_sym = self.ctx.cls.info.get(field.name)
+                if field_sym is not None and field_sym.node is not None:
+                    error_context = field_sym.node
+                self.api.fail(
+                    f"Cannot find model {field.related_model!r} referenced in field {field.name!r}",
+                    ctx=error_context,
+                )
+                self.add_new_node_to_model_class(field.attname, AnyType(TypeOfAny.explicit))
+                continue
 
-                if related_model_cls._meta.abstract:
-                    continue
+            if related_model_cls._meta.abstract:
+                continue
 
-                rel_target_field = self.django_context.get_related_target_field(related_model_cls, field)
-                if not rel_target_field:
-                    continue
+            rel_target_field = self.django_context.get_related_target_field(related_model_cls, field)
+            if not rel_target_field:
+                continue
 
-                try:
-                    field_info = self.lookup_class_typeinfo_or_incomplete_defn_error(rel_target_field.__class__)
-                except helpers.IncompleteDefnException as exc:
-                    if not self.api.final_iteration:
-                        raise exc
-                    else:
-                        continue
+            try:
+                field_info = self.lookup_class_typeinfo_or_incomplete_defn_error(rel_target_field.__class__)
+            except helpers.IncompleteDefnException as exc:
+                if not self.api.final_iteration:
+                    raise exc
+                else:
+                    continue
 
-                is_nullable = self.django_context.get_field_nullability(field, None)
-                set_type, get_type = get_field_descriptor_types(
-                    field_info, is_set_nullable=is_nullable, is_get_nullable=is_nullable
-                )
-                self.add_new_node_to_model_class(field.attname, Instance(field_info, [set_type, get_type]))
+            is_nullable = self.django_context.get_field_nullability(field, None)
+            set_type, get_type = get_field_descriptor_types(
+                field_info, is_set_nullable=is_nullable, is_get_nullable=is_nullable
+            )
+            self.add_new_node_to_model_class(field.attname, Instance(field_info, [set_type, get_type]))
 
 
 class AddManagers(ModelClassInitializer):