diff --git a/django-stubs/db/models/fields/related.pyi b/django-stubs/db/models/fields/related.pyi index 3d872721c..8ed1954a8 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 d30627438..2c041994b 100644 --- a/mypy_django_plugin/django/context.py +++ b/mypy_django_plugin/django/context.py @@ -143,6 +143,20 @@ def get_field_lookup_exact_type( 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 a4732a325..2dd2584cd 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -253,9 +253,12 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: 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 b3b12dcf2..60cc66e75 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