pagination max_depth (Sourcery refactored)

docs/custom_queryset.md

@@ -6,7 +6,8 @@ flask-mongoengine attaches the following methods to Mongoengine's default QueryS
   Optional arguments: *message* - custom message to display.
 * **first_or_404**: same as above, except for .first().
   Optional arguments: *message* - custom message to display.
-* **paginate**: paginates the QuerySet. Takes two arguments, *page* and *per_page*.
+* **paginate**: paginates the QuerySet. Takes two required arguments, *page* and *per_page*.
+   And one optional arguments *max_depth*.
 * **paginate_field**: paginates a field from one document in the QuerySet.
   Arguments: *field_name*, *doc_id*, *page*, *per_page*.
 

flask_mongoengine/pagination.py

@@ -8,7 +8,13 @@
 
 
 class Pagination(object):
-    def __init__(self, iterable, page, per_page):
+    def __init__(self, iterable, page: int, per_page: int, max_depth: int = None):
+        """
+        :param iterable: iterable object .
+        :param page: Required page number start from 1.
+        :param per_page: Required number of documents per page.
+        :param max_depth: Option for limit number of dereference documents.
+        """
 
         if page < 1:
             abort(404)
@@ -19,11 +25,9 @@ def __init__(self, iterable, page, per_page):
 
         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()
-            )
+            self.items = self.iterable.skip(self.per_page * (self.page - 1)).limit(self.per_page)
+            if max_depth is not None:
+                self.items = self.items.select_related(max_depth)
         else:
             start_index = (page - 1) * per_page
             end_index = page * per_page

tests/test_pagination.py

@@ -1,14 +1,30 @@
+import flask
 import pytest
 from werkzeug.exceptions import NotFound
 
 from flask_mongoengine import ListFieldPagination, Pagination
 
 
-def test_queryset_paginator(app, todo):
+@pytest.fixture(autouse=True)
+def setup_endpoints(app, todo):
     Todo = todo
     for i in range(42):
         Todo(title=f"post: {i}").save()
 
+    @app.route("/")
+    def index():
+        page = int(flask.request.form.get("page"))
+        per_page = int(flask.request.form.get("per_page"))
+        query_set = Todo.objects().paginate(page=page, per_page=per_page)
+        return {
+            'data': list(query_set.items),
+            'total': query_set.total,
+            'has_next': query_set.has_next,
+        }
+
+
+def test_queryset_paginator(app, todo):
+    Todo = todo
     with pytest.raises(NotFound):
         Pagination(iterable=Todo.objects, page=0, per_page=10)
 
@@ -90,3 +106,26 @@ def _test_paginator(paginator):
         # Paginate to the next page
         if i < 5:
             paginator = paginator.next()
+
+
+
+def test_flask_pagination(app, todo):
+    client = app.test_client()
+    response = client.get("/", data={"page": 0, "per_page": 10})
+    print(response.status_code)
+    assert response.status_code == 404
+
+    response = client.get("/", data={"page": 6, "per_page": 10})
+    print(response.status_code)
+    assert response.status_code == 404
+
+
+def test_flask_pagination_next(app, todo):
+    client = app.test_client()
+    has_next = True
+    page = 1
+    while has_next:
+        response = client.get("/", data={"page": page, "per_page": 10})
+        assert response.status_code == 200
+        has_next = response.json['has_next']
+        page += 1