Fix Field.get_attribute() return and RelatedField.to_representation() parameter for PKOnlyObject (#852)

Looking at the implementation of `Serializer.to_representation()`, the value returned by `Field.get_attribute()` is passed to `Field.to_representation()`, whose argument is declared to be `_VT`. Thus the declaration of `Field.get_attribute()` should be changed.

Adjusts the typing for relational serializer fields so they return the actual model instances or the lightweight PKOnlyObject Django REST Framework sometimes uses, instead of the already-serialized values. This matches DRF’s real behavior: get_attribute hands to_representation either the full object or the PK-only placeholder, so the new annotations describe exactly what callers receive and avoid type errors when the optimization kicks in.

Signed-off-by: Emmanuel Ferdman <[email protected]>

Author: emmanuel-ferdman

rest_framework-stubs/fields.pyi

@@ -113,7 +113,7 @@ class Field(Generic[_VT, _DT, _RP, _IN]):
     def get_validators(self) -> list[Validator[_VT]]: ...
     def get_initial(self) -> _VT | None: ...
     def get_value(self, dictionary: Mapping[Any, Any]) -> Any: ...
-    def get_attribute(self, instance: _IN) -> _RP | None: ...
+    def get_attribute(self, instance: _IN) -> _VT | None: ...
     def get_default(self) -> _VT | None: ...
     def validate_empty_values(self, data: Any) -> tuple[bool, Any]: ...
     def run_validation(self, data: Any = ...) -> Any: ...

rest_framework-stubs/relations.pyi

@@ -68,7 +68,8 @@ class RelatedField(Field[_MT, _DT, _PT, Any]):
     @property
     def grouped_choices(self) -> dict: ...
     def iter_options(self) -> Iterable[Option]: ...
-    def get_attribute(self, instance: _MT) -> _PT | None: ...
+    def get_attribute(self, instance: _MT) -> _MT | PKOnlyObject | None: ...  # type: ignore[override]
+    def to_representation(self, value: _MT | PKOnlyObject) -> _PT: ...
     def display_value(self, instance: _MT) -> str: ...
 
 class StringRelatedField(RelatedField[_MT, _MT, str]): ...
@@ -160,7 +161,7 @@ class SlugRelatedField(RelatedField[_MT, str, str]):
         style: dict[str, str] | None = ...,
     ) -> None: ...
     def to_internal_value(self, data: Any) -> _MT: ...
-    def to_representation(self, obj: _MT) -> str: ...
+    def to_representation(self, obj: _MT | PKOnlyObject) -> str: ...
 
 class ManyRelatedField(Field[Sequence[Any], Sequence[Any], list[Any], Any]):
     default_empty_html: list[object]

tests/typecheck/test_fields.yml

@@ -57,6 +57,14 @@
         x: str | None = CharField().get_initial()
         y: int | None = CharField().get_initial()  # E: Incompatible types in assignment (expression has type "str | None", variable has type "int | None")  [assignment]
 
+-   case: field_get_attribute_returns_value_type
+    main: |
+        from rest_framework import serializers
+
+        field = serializers.CharField()
+        result = field.get_attribute(object())
+        reveal_type(result)  # N: Revealed type is "builtins.str | None"
+
 -   case: float_field_args_fields
     main: |
       from rest_framework.fields import FloatField

tests/typecheck/test_serializers.yml

@@ -49,6 +49,28 @@
         class Meta:
             model = User
 
+- case: related_field_get_attribute_model_or_pk
+  main: |
+    from typing import cast
+    from django.contrib.auth.models import User
+    from rest_framework import serializers
+    from rest_framework.relations import PKOnlyObject
+
+    field = serializers.PrimaryKeyRelatedField(queryset=User.objects.all())
+    user = cast(User, object())
+    value = field.get_attribute(user)
+    reveal_type(value)  # N: Revealed type is "django.contrib.auth.models.User | rest_framework.relations.PKOnlyObject | None"
+
+- case: slug_related_field_accepts_pk_object
+  main: |
+    from django.contrib.auth.models import User
+    from rest_framework import serializers
+    from rest_framework.relations import PKOnlyObject
+
+    field = serializers.SlugRelatedField(slug_field="username", queryset=User.objects.all())
+    obj = PKOnlyObject(pk=1)
+    reveal_type(field.to_representation(obj))  # N: Revealed type is "builtins.str"
+
 - case: test_hyperlinked_model_serializer_with_customized_serializer_field_mapping
   main: |
     from rest_framework import serializers