Merge pull request #485 from idoshr/pagigation

Improved: Pagination performance

flask_mongoengine/pagination.py

@@ -16,14 +16,20 @@ def __init__(self, iterable, page, per_page):
         self.iterable = iterable
         self.page = page
         self.per_page = per_page
-        self.total = len(iterable)
 
-        start_index = (page - 1) * per_page
-        end_index = page * per_page
-
-        self.items = iterable[start_index:end_index]
-        if isinstance(self.items, QuerySet):
-            self.items = self.items.select_related()
+        if isinstance(self.iterable, QuerySet):
+            self.total = iterable.count()
+            self.items = (
+                self.iterable.skip(self.per_page * (self.page - 1))
+                .limit(self.per_page)
+                .select_related()
+            )
+        else:
+            start_index = (page - 1) * per_page
+            end_index = page * per_page
+
+            self.total = len(iterable)
+            self.items = iterable[start_index:end_index]
         if not self.items and page != 1:
             abort(404)
 

tests/test_pagination.py

@@ -18,6 +18,12 @@ def test_queryset_paginator(app, todo):
     paginator = Pagination(Todo.objects, 1, 10)
     _test_paginator(paginator)
 
+    for page in range(1, 10):
+        for index, todo in enumerate(
+            Todo.objects.paginate(page=page, per_page=5).items
+        ):
+            assert todo.title == f"post: {(page-1) * 5 + index}"
+
 
 def test_paginate_plain_list():
     with pytest.raises(NotFound):