Fix related_model lookups for un-swapped models

typeddjango · May 12, 2023 · 228e285 · 228e285
1 parent 62ed7df
commit 228e285
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py
@@ -128,7 +128,7 @@ def get_field_lookup_exact_type(
         self, api: TypeChecker, field: Union["Field[Any, Any]", ForeignObjectRel]
     ) -> MypyType:
         if isinstance(field, (RelatedField, ForeignObjectRel)):
-            related_model_cls = field.related_model
+            related_model_cls = self.get_field_related_model_cls(field)
             primary_key_field = self.get_primary_key_field(related_model_cls)
             primary_key_type = self.get_field_get_type(api, primary_key_field, method="init")
 
@@ -364,13 +364,13 @@ def _resolve_field_from_parts(
 
             field = currently_observed_model._meta.get_field(field_part)
             if isinstance(field, RelatedField):
-                currently_observed_model = field.related_model
+                currently_observed_model = self.get_field_related_model_cls(field)
                 model_name = currently_observed_model._meta.model_name
                 if model_name is not None and field_part == (model_name + "_id"):
                     field = self.get_primary_key_field(currently_observed_model)
 
             if isinstance(field, ForeignObjectRel):
-                currently_observed_model = field.related_model
+                currently_observed_model = self.get_field_related_model_cls(field)
 
         # Guaranteed by `query.solve_lookup_type` before.
         assert isinstance(field, (Field, ForeignObjectRel))