From db14454199c8634443881d934689e374c80a4810 Mon Sep 17 00:00:00 2001
From: Marti Raudsepp <marti@juffo.org>
Date: Fri, 30 Sep 2022 10:05:57 +0300
Subject: [PATCH] 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 |  2 ++
 mypy_django_plugin/django/context.py      | 14 +++++++++++++
 mypy_django_plugin/transformers/models.py |  7 +++++--
 tests/typecheck/fields/test_related.yml   | 24 +++++++++++++++++++++++
 4 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/django-stubs/db/models/fields/related.pyi b/django-stubs/db/models/fields/related.pyi
index 3d87272..8ed1954 100644
--- a/django-stubs/db/models/fields/related.pyi
+++ b/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,
diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py
index d306274..2c04199 100644
--- a/mypy_django_plugin/django/context.py
+++ b/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):
diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py
index a4732a3..2dd2584 100644
--- a/mypy_django_plugin/transformers/models.py
+++ b/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
diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml
index b3b12dc..60cc66e 100644
--- a/tests/typecheck/fields/test_related.yml
+++ b/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