Fix #144 -- Improve DRM serialization performance by filtering source and ratio (#148)

Author: amureki

README.md

@@ -246,9 +246,21 @@ class PictureSerializer(serializers.Serializer):
     picture = PictureField()
 ```
 
-You may provide optional GET parameters to the serializer, to specify the aspect
-ratio and breakpoints you want to include in the response. The parameters are
-prefixed with the `fieldname_` to avoid conflicts with other fields.
+The response can be restricted to a single aspect ratio and image source, by
+providing the `aspect_ratio` and `image_source` arguments to the field.
+
+```python
+from rest_framework import serializers
+from pictures.contrib.rest_framework import PictureField
+
+class PictureSerializer(serializers.Serializer):
+    picture = PictureField(aspect_ratio="16/9", image_source="WEBP")
+```
+
+You also may provide optional GET parameters to the serializer,
+to specify the aspect ratio and breakpoints you want to include in the response.
+The parameters are prefixed with the `fieldname_`
+to avoid conflicts with other fields.
 
 ```bash
 curl http://localhost:8000/api/path/?picture_ratio=16%2F9&picture_m=6&picture_l=4

pictures/contrib/rest_framework.py

@@ -19,23 +19,52 @@ def default(obj):
 class PictureField(serializers.ReadOnlyField):
     """Read-only field for all aspect ratios and sizes of the image."""
 
+    def __init__(self, aspect_ratio=None, image_source=None, **kwargs):
+        self.aspect_ratio = aspect_ratio
+        self.image_source = image_source
+        super().__init__(**kwargs)
+
     def to_representation(self, obj: PictureFieldFile):
         if not obj:
             return None
-        payload = {
+
+        base_payload = {
             "url": obj.url,
             "width": obj.width,
             "height": obj.height,
-            "ratios": {
-                ratio: {
-                    "sources": {
-                        f"image/{file_type.lower()}": sizes
-                        for file_type, sizes in sources.items()
-                    },
-                }
-                for ratio, sources in obj.aspect_ratios.items()
-            },
         }
+
+        # if aspect_ratio is set, only return that aspect ratio to reduce payload size
+        if self.aspect_ratio and self.image_source:
+            try:
+                sizes = obj.aspect_ratios[self.aspect_ratio][self.image_source]
+            except KeyError as e:
+                raise ValueError(
+                    f"Invalid ratio {self.aspect_ratio} or image source {self.image_source}. Choices are: {', '.join(filter(None, obj.aspect_ratios.keys()))}"
+                ) from e
+            payload = {
+                **base_payload,
+                "ratios": {
+                    self.aspect_ratio: {
+                        "sources": {f"image/{self.image_source.lower()}": sizes}
+                    }
+                },
+            }
+        else:
+            payload = {
+                **base_payload,
+                "ratios": {
+                    ratio: {
+                        "sources": {
+                            f"image/{file_type.lower()}": sizes
+                            for file_type, sizes in sources.items()
+                        },
+                    }
+                    for ratio, sources in obj.aspect_ratios.items()
+                },
+            }
+
+        # if the request has query parameters, filter the payload
         try:
             query_params: QueryDict = self.context["request"].GET
         except KeyError:

tests/contrib/test_rest_framework.py

@@ -12,10 +12,23 @@
 
 class ProfileSerializer(serializers.ModelSerializer):
     image = rest_framework.PictureField(source="picture")
+    image_mobile = rest_framework.PictureField(
+        source="picture", aspect_ratio="3/2", image_source="WEBP"
+    )
+
+    class Meta:
+        model = models.Profile
+        fields = ["image", "image_mobile"]
+
+
+class ProfileSerializerWithInvalidData(serializers.ModelSerializer):
+    image_invalid = rest_framework.PictureField(
+        source="picture", aspect_ratio="21/11", image_source="GIF"
+    )
 
     class Meta:
         model = models.Profile
-        fields = ["image"]
+        fields = ["image_invalid"]
 
 
 def test_default(settings):
@@ -264,3 +277,51 @@ def test_to_representation__with_false_str_container(
         with pytest.raises(ValueError) as e:
             serializer.data["image"]
         assert str(e.value) == "Container width is not a number: not_a_number"
+
+    @pytest.mark.django_db
+    def test_to_representation__with_prefiltered_aspect_ratio_and_source(
+        self, image_upload_file, settings
+    ):
+        settings.PICTURES["USE_PLACEHOLDERS"] = False
+
+        profile = models.Profile.objects.create(picture=image_upload_file)
+        serializer = ProfileSerializer(profile)
+
+        assert serializer.data["image_mobile"] == {
+            "url": "/media/testapp/profile/image.png",
+            "width": 800,
+            "height": 800,
+            "ratios": {
+                "3/2": {
+                    "sources": {
+                        "image/webp": {
+                            "800": "/media/testapp/profile/image/3_2/800w.webp",
+                            "100": "/media/testapp/profile/image/3_2/100w.webp",
+                            "200": "/media/testapp/profile/image/3_2/200w.webp",
+                            "300": "/media/testapp/profile/image/3_2/300w.webp",
+                            "400": "/media/testapp/profile/image/3_2/400w.webp",
+                            "500": "/media/testapp/profile/image/3_2/500w.webp",
+                            "600": "/media/testapp/profile/image/3_2/600w.webp",
+                            "700": "/media/testapp/profile/image/3_2/700w.webp",
+                        }
+                    }
+                }
+            },
+        }
+
+    @pytest.mark.django_db
+    def test_to_representation__with_prefiltered_aspect_ratio_and_source__raise_value_error(
+        self, image_upload_file, settings
+    ):
+        settings.PICTURES["USE_PLACEHOLDERS"] = False
+
+        profile = models.Profile.objects.create(picture=image_upload_file)
+
+        serializer = ProfileSerializerWithInvalidData(profile)
+        with pytest.raises(ValueError) as e:
+            serializer.data["image_invalid"]
+
+        assert (
+            str(e.value)
+            == "Invalid ratio 21/11 or image source GIF. Choices are: 1/1, 3/2, 16/9"
+        )