allow Prefetch objects to be used for prefetching

Author: tonisvain

.gitignore

@@ -4,6 +4,7 @@ db.sqlite3
 *.py[cod]
 __pycache__
 env
+venv
 .python-version
 
 # C extensions

drf_sideloading/mixins.py

@@ -5,6 +5,7 @@
 
 from itertools import chain
 
+from django.db.models import Prefetch
 from rest_framework.renderers import BrowsableAPIRenderer
 from rest_framework.response import Response
 from rest_framework.serializers import ListSerializer
@@ -90,14 +91,11 @@ def get_sideloading_prefetches(self):
         cleaned_prefetches = {}
         for k, v in prefetches.items():
             if v is not None:
-                if isinstance(v, list):
-                    cleaned_prefetches[k] = v
-                elif isinstance(v, six.string_types):
-                    cleaned_prefetches[k] = [v]
-                else:
-                    raise RuntimeError(
-                        "Sideloadable prefetch values must be presented either as a list or a string"
-                    )
+                if not isinstance(v, list):
+                    v = [v]
+                for vi in v:
+                    if not isinstance(vi, (six.string_types, Prefetch)):
+                        raise RuntimeError("Sideloadable prefetch values must be a list of strings or Prefetch objects")
         return cleaned_prefetches
 
     def initialize_request(self, request, *args, **kwargs):

tests/test_products_api.py

@@ -3,6 +3,7 @@
 
 from __future__ import unicode_literals
 
+from django.db.models import Prefetch
 from rest_framework.permissions import BasePermission
 from rest_framework.renderers import BrowsableAPIRenderer, JSONRenderer
 from rest_framework.settings import api_settings
@@ -388,19 +389,14 @@ class Meta:
         ProductViewSet.sideloading_serializer_class = ProductSideloadableSerializer
 
     def test_correct_exception_raised(self):
-        with self.assertRaises(RuntimeError) as cm:
+        expected_error_message = "Sideloadable prefetch values must be a list of strings or Prefetch objects"
+        with self.assertRaisesMessage(RuntimeError, expected_error_message):
             self.client.get(
                 reverse("product-list"),
                 data={"sideload": "categories,suppliers,partners"},
                 format="json",
             )
 
-        expected_error_message = "Sideloadable prefetch values must be presented either as a list or a string"
-
-        raised_exception = cm.exception
-        self.assertEqual(str(raised_exception), expected_error_message)
-
-
 class TestDrfSideloadingValidPrefetches(TestCase):
     """Run tests while including mixin but not defining sideloading"""
 
@@ -412,13 +408,15 @@ class ProductSideloadableSerializer(SideLoadableSerializer):
             products = ProductSerializer(many=True)
             categories = CategorySerializer(source="category", many=True)
             suppliers = SupplierSerializer(source="supplier", many=True)
+            suppliers_ordered_by_name = SupplierSerializer(source="supplier", many=True)
             partners = PartnerSerializer(many=True)
 
             class Meta:
                 primary = "products"
                 prefetches = {
                     "categories": "category",
                     "suppliers": ["supplier"],
+                    "suppliers_ordered_by_name": Prefetch("supplier", queryset=Supplier.objects.order_by("name")),
                     "partners": None,
                 }
 
@@ -427,7 +425,7 @@ class Meta:
     def test_sideloading_with_prefetches(self):
         self.client.get(
             reverse("product-list"),
-            data={"sideload": "categories,suppliers,partners"},
+            data={"sideload": "categories,suppliers,suppliers_ordered_by_name,partners"},
             format="json",
         )