Fix pk-only optimization for properties (encode#7142)

* Add callable/prop tests for pk-only optimization * Fix related field pk-only optimization for props

Author: rpkilby

rest_framework/relations.py

@@ -175,8 +175,13 @@ def get_attribute(self, instance):
  value = attribute_instance.serializable_value(self.source_attrs[-1])
  if is_simple_callable(value):
  # Handle edge case where the relationship `source` argument
- # points to a `get_relationship()` method on the model
- value = value().pk
+ # points to a `get_relationship()` method on the model.
+ value = value()
+
+ # Handle edge case where relationship `source` argument points
+ # to an instance instead of a pk (e.g., a `@property`).
+ value = getattr(value, 'pk', value)
+
  return PKOnlyObject(pk=value)
  except AttributeError:
  pass

tests/models.py

@@ -37,6 +37,15 @@ class ManyToManySource(RESTFrameworkModel):
 class ForeignKeyTarget(RESTFrameworkModel):
  name = models.CharField(max_length=100)
 
+ def get_first_source(self):
+ """Used for testing related field against a callable."""
+ return self.sources.all().order_by('pk')[0]
+
+ @property
+ def first_source(self):
+ """Used for testing related field against a property."""
+ return self.sources.all().order_by('pk')[0]
+
 
 class UUIDForeignKeyTarget(RESTFrameworkModel):
  uuid = models.UUIDField(primary_key=True, default=uuid.uuid4)

tests/test_relations_pk.py

@@ -30,6 +30,25 @@ class Meta:
  fields = ('id', 'name', 'sources')
 
 
+class ForeignKeyTargetCallableSourceSerializer(serializers.ModelSerializer):
+ first_source = serializers.PrimaryKeyRelatedField(
+ source='get_first_source',
+ read_only=True,
+ )
+
+ class Meta:
+ model = ForeignKeyTarget
+ fields = ('id', 'name', 'first_source')
+
+
+class ForeignKeyTargetPropertySourceSerializer(serializers.ModelSerializer):
+ first_source = serializers.PrimaryKeyRelatedField(read_only=True)
+
+ class Meta:
+ model = ForeignKeyTarget
+ fields = ('id', 'name', 'first_source')
+
+
 class ForeignKeySourceSerializer(serializers.ModelSerializer):
  class Meta:
  model = ForeignKeySource
@@ -389,6 +408,34 @@ class Meta:
  assert len(queryset) == 1
 
 
+class PKRelationTests(TestCase):
+
+ def setUp(self):
+ self.target = ForeignKeyTarget.objects.create(name='target-1')
+ ForeignKeySource.objects.create(name='source-1', target=self.target)
+ ForeignKeySource.objects.create(name='source-2', target=self.target)
+
+ def test_relation_field_callable_source(self):
+ serializer = ForeignKeyTargetCallableSourceSerializer(self.target)
+ expected = {
+ 'id': 1,
+ 'name': 'target-1',
+ 'first_source': 1,
+ }
+ with self.assertNumQueries(1):
+ self.assertEqual(serializer.data, expected)
+
+ def test_relation_field_property_source(self):
+ serializer = ForeignKeyTargetPropertySourceSerializer(self.target)
+ expected = {
+ 'id': 1,
+ 'name': 'target-1',
+ 'first_source': 1,
+ }
+ with self.assertNumQueries(1):
+ self.assertEqual(serializer.data, expected)
+
+
 class PKNullableForeignKeyTests(TestCase):
  def setUp(self):
  target = ForeignKeyTarget(name='target-1')