Fix type of <fieldname>_id when using ForeignKey(to_field=) (#1176)

* Fix type of <fieldname>_id when using ForeignKey(to_field=)

Previously mypy_django_plugin would always use the field type of target
model's primary key, but `to_field` can refer to a different field type.

* Fixes

* More fixes

django-stubs/db/models/fields/related.pyi

@@ -61,6 +61,8 @@ class RelatedField(FieldCacheMixin, Field[_ST, _GT]):
 class ForeignObject(RelatedField[_ST, _GT]):
     remote_field: ForeignObjectRel
     rel_class: Type[ForeignObjectRel]
+    from_fields: Sequence[str]
+    to_fields: Sequence[str | None]  # None occurs in ForeignKey, where to_field defaults to None
     swappable: bool
     def __init__(
         self,

mypy_django_plugin/django/context.py

@@ -143,6 +143,20 @@ class DjangoContext:
             return AnyType(TypeOfAny.explicit)
         return helpers.get_private_descriptor_type(field_info, "_pyi_lookup_exact_type", is_nullable=field.null)
 
+    def get_related_target_field(
+        self, related_model_cls: Type[Model], field: "ForeignKey[Any, Any]"
+    ) -> "Optional[Field[Any, Any]]":
+        # ForeginKey only supports one `to_fields` item (ForeignObject supports many)
+        assert len(field.to_fields) == 1
+        to_field_name = field.to_fields[0]
+        if to_field_name:
+            rel_field = related_model_cls._meta.get_field(to_field_name)
+            if not isinstance(rel_field, Field):
+                return None  # Not supported
+            return rel_field
+        else:
+            return self.get_primary_key_field(related_model_cls)
+
     def get_primary_key_field(self, model_cls: Type[Model]) -> "Field[Any, Any]":
         for field in model_cls._meta.get_fields():
             if isinstance(field, Field):

mypy_django_plugin/transformers/models.py

@@ -253,9 +253,12 @@ class AddRelatedModelsId(ModelClassInitializer):
                 if related_model_cls._meta.abstract:
                     continue
 
-                rel_primary_key_field = self.django_context.get_primary_key_field(related_model_cls)
+                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_primary_key_field.__class__)
+                    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

tests/typecheck/fields/test_related.yml

@@ -38,6 +38,30 @@
                     publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE)
                     owner = models.ForeignKey(db_column='model_id', to='auth.User', on_delete=models.CASCADE)
 
+-   case: foreign_key_field_custom_to_field
+    main: |
+        from myapp.models import Book, Publisher
+        from uuid import UUID
+        book = Book()
+        book.publisher = Publisher()
+        reveal_type(book.publisher_id)  # N: Revealed type is "uuid.UUID"
+        book.publisher_id = '821850bb-c105-426f-b340-3974419d00ca'
+        book.publisher_id = UUID('821850bb-c105-426f-b340-3974419d00ca')
+        book.publisher_id = [1]  # E: Incompatible types in assignment (expression has type "List[int]", variable has type "Union[str, UUID]")
+        book.publisher_id = Publisher()  # E: Incompatible types in assignment (expression has type "Publisher", variable has type "Union[str, UUID]")
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class Publisher(models.Model):
+                    id = models.BigAutoField(primary_key=True)
+                    uuid = models.UUIDField(unique=True)
+                class Book(models.Model):
+                    publisher = models.ForeignKey(to=Publisher, to_field='uuid', on_delete=models.CASCADE)
+
 -   case: foreign_key_field_different_order_of_params
     main: |
         from myapp.models import Book, Publisher