Improve type checking in pagination.py

Author: martin056

styleguide_example/api/pagination.py

@@ -1,30 +1,33 @@
 from collections import OrderedDict
+from typing import Optional, Sequence, Tuple, Type, Union
 from urllib.parse import parse_qs, urlparse
 
-from rest_framework.pagination import LimitOffsetPagination as _LimitOffsetPagination
+from django.db.models import QuerySet
+from rest_framework import serializers
+from rest_framework.pagination import BasePagination
 from rest_framework.pagination import CursorPagination as _CursorPagination
+from rest_framework.pagination import LimitOffsetPagination as _LimitOffsetPagination
+from rest_framework.request import Request
 from rest_framework.response import Response
+from rest_framework.views import APIView
 
 
-def get_paginated_response(*, pagination_class, serializer_class, queryset, request, view):
-    paginator = pagination_class()
+class _TurnOffPaginationSerializer(serializers.Serializer):
+    paginate = serializers.BooleanField(default=True)
 
-    page = paginator.paginate_queryset(queryset, request, view=view)
 
-    if page is not None:
-        serializer = serializer_class(page, many=True)
-        return paginator.get_paginated_response(serializer.data)
-
-    serializer = serializer_class(queryset, many=True)
+def turn_off_pagination(data):
+    serializer = _TurnOffPaginationSerializer(data=data)
+    serializer.is_valid(raise_exception=True)
 
-    return Response(data=serializer.data)
+    return serializer.validated_data["paginate"]
 
 
 class LimitOffsetPagination(_LimitOffsetPagination):
     default_limit = 10
     max_limit = 50
 
-    def get_count(self, queryset) -> int:
+    def get_count(self, queryset: Union[QuerySet, Sequence]) -> int:
         """
         Determine an object count, supporting either querysets or regular lists.
         """
@@ -42,44 +45,52 @@ def get_count(self, queryset) -> int:
             return len(queryset)
 
     def get_paginated_data(self, data):
-        return OrderedDict([
-            ('limit', self.limit),
-            ('offset', self.offset),
-            ('count', self.count),
-            ('next', self.get_next_link()),
-            ('previous', self.get_previous_link()),
-            ('results', data)
-        ])
+        return OrderedDict(
+            [
+                ("limit", self.limit),
+                ("offset", self.offset),
+                ("count", self.count),
+                ("next", self.get_next_link()),
+                ("previous", self.get_previous_link()),
+                ("results", data),
+            ]
+        )
 
     def get_paginated_response(self, data):
         """
         We redefine this method in order to return `limit` and `offset`.
         This is used by the frontend to construct the pagination itself.
         """
-        return Response(OrderedDict([
-            ('limit', self.limit),
-            ('offset', self.offset),
-            ('count', self.count),
-            ('next', self.get_next_link()),
-            ('previous', self.get_previous_link()),
-            ('results', data)
-        ]))
+        return Response(
+            OrderedDict(
+                [
+                    ("limit", self.limit),
+                    ("offset", self.offset),
+                    ("count", self.count),
+                    ("next", self.get_next_link()),
+                    ("previous", self.get_previous_link()),
+                    ("results", data),
+                ]
+            )
+        )
 
 
 class CursorPagination(_CursorPagination):
     page_size = 50  # Return 50 items by default
 
-    def __init__(self, ordering):
+    def __init__(self, ordering: Optional[str]):
         self.ordering: str = ordering or "-created_at"
 
-    def get_ordering(self, request, queryset, view):
+    def get_ordering(
+        self, request: Request, queryset: QuerySet, view: APIView
+    ) -> Tuple[str]:
         # The DRF CursorPagination expects the ordering as a tuple
         if isinstance(self.ordering, str):
             return (self.ordering,)
 
         return tuple(self.ordering)
 
-    def _get_cursor(self, url):
+    def _get_cursor(self, url: Optional[str]) -> Optional[str]:
         if not url:
             return None
 
@@ -109,3 +120,42 @@ def get_paginated_response(self, data):
                 ]
             )
         )
+
+
+def _init_pagination_class(
+    pagination_class: Type[BasePagination],
+    ordering: Optional[str],
+) -> BasePagination:
+    if isinstance(pagination_class, CursorPagination):
+        return pagination_class(ordering=ordering)
+
+    return pagination_class()
+
+
+def response_paginate(
+    *,
+    pagination_class: Type[BasePagination],
+    serializer_class: Type[serializers.Serializer],
+    queryset: QuerySet,
+    request: Request,
+    view: APIView,
+    ordering: Optional[str] = "-created_at"
+) -> Response:
+    paginate = turn_off_pagination(data=request.GET)
+
+    if not paginate:
+        data = serializer_class(queryset, many=True).data
+
+        return Response(data=data)
+
+    paginator = _init_pagination_class(pagination_class, ordering)
+
+    page = paginator.paginate_queryset(queryset, request, view=view)
+
+    if page is not None:
+        serializer = serializer_class(page, many=True)
+        return paginator.get_paginated_response(serializer.data)
+
+    serializer = serializer_class(queryset, many=True)
+
+    return Response(data=serializer.data)

styleguide_example/api/tests/pagination/test_get_paginated_response.py styleguide_example/api/tests/pagination/test_response_paginate.py

@@ -6,7 +6,7 @@
 from rest_framework.views import APIView
 from rest_framework import serializers
 
-from styleguide_example.api.pagination import get_paginated_response, LimitOffsetPagination
+from styleguide_example.api.pagination import response_paginate, LimitOffsetPagination
 
 from styleguide_example.users.services import user_create
 from styleguide_example.users.models import BaseUser
@@ -24,7 +24,7 @@ class Meta:
     def get(self, request):
         queryset = BaseUser.objects.order_by('id')
 
-        response = get_paginated_response(
+        response = response_paginate(
             pagination_class=self.Pagination,
             serializer_class=self.OutputSerializer,
             queryset=queryset,
@@ -35,7 +35,7 @@ def get(self, request):
         return response
 
 
-class GetPaginatedResponseTests(TestCase):
+class ResponsePaginateTests(TestCase):
     def setUp(self):
         self.factory = APIRequestFactory()
 

styleguide_example/users/apis.py

@@ -1,7 +1,7 @@
 from rest_framework.views import APIView
 from rest_framework import serializers
 
-from styleguide_example.api.pagination import get_paginated_response, LimitOffsetPagination
+from styleguide_example.api.pagination import response_paginate, LimitOffsetPagination
 
 from styleguide_example.users.selectors import user_list
 from styleguide_example.users.models import BaseUser
@@ -34,7 +34,7 @@ def get(self, request):
 
         users = user_list(filters=filters_serializer.validated_data)
 
-        return get_paginated_response(
+        return response_paginate(
             pagination_class=self.Pagination,
             serializer_class=self.OutputSerializer,
             queryset=users,