Implement pagination headers

Fixes #19

flask_stupe/json.py

@@ -124,6 +124,10 @@ def make_response(self, rv):
 
         rv = jsonify(rv)
         rv.status_code = code
+
+        if request.response_headers:
+            rv.headers.extend(request.response_headers)
+
         return rv
 
     def __init__(self, *args, **kwargs):

flask_stupe/pagination.py

@@ -8,10 +8,35 @@
 
 
 if pymongo:
-    def _paginate(cursor, skip=None, limit=None, sort=None, count=True):
-        metadata = getattr(request, "metadata", None)
-        if count and isinstance(metadata, dict):
-            metadata.update(count=cursor.count())
+
+    def _get_link_header(limit, skip, total_count=None):
+        template = '<{}?limit={}&skip={{skip}}>; rel="{{rel}}"'
+        template = template.format(request.base_url, limit)
+        links = [
+            template.format(skip=0, rel="first"),
+            template.format(skip=skip - limit, rel="prev"),
+            template.format(skip=skip + limit, rel="next")
+        ]
+        if total_count:
+            links.append(template.format(skip=total_count - limit, rel="last"))
+        return ", ".join(links)
+
+    def _paginate(cursor, skip=None, limit=None, sort=None, count=True,
+                  headers=False):
+        if headers:
+            headers = getattr(request, "response_headers", None)
+            if isinstance(headers, dict):
+                total_count = None
+                if count:
+                    total_count = cursor.count()
+                    headers["X-Total-Count"] = total_count
+                if limit is not None and skip is not None:
+                    headers["Link"] = _get_link_header(limit, skip,
+                                                       total_count)
+        else:
+            metadata = getattr(request, "metadata", None)
+            if count and isinstance(metadata, dict):
+                metadata.update(count=cursor.count())
 
         skip = request.args.get("skip", skip, type=int)
         if skip is not None:
@@ -35,16 +60,17 @@ def _paginate(cursor, skip=None, limit=None, sort=None, count=True):
         return cursor
 
     def paginate(function_or_cursor=None, skip=None, limit=None, sort=None,
-                 count=True):
+                 count=True, headers=False):
         """Apply pagination to the given MongoDB cursor or function"""
         if isinstance(function_or_cursor, pymongo.cursor.Cursor):
-            return _paginate(function_or_cursor, skip, limit, sort, count)
+            return _paginate(function_or_cursor, skip, limit, sort, count,
+                             headers)
 
         def __decorator(function):
             @functools.wraps(function)
             def __wrapper(*args, **kwargs):
                 cursor = function(*args, **kwargs)
-                return _paginate(cursor, skip, limit, sort, count)
+                return _paginate(cursor, skip, limit, sort, count, headers)
             return __wrapper
 
         if function_or_cursor:

flask_stupe/request.py

@@ -8,6 +8,7 @@ def __init__(self, *args, **kwargs):
 
         #: Store additionnal data about the request.
         self.metadata = {}
+        self.response_headers = {}
 
 
 __all__ = ["Request"]

tests/json.py

@@ -182,6 +182,21 @@ def foo():
     assert response_dict["count"] == 3
 
 
+def test_stupeflask_response_with_paginate_headers(json_app, client):
+    encoder_rules.append((Cursor, lambda c: c.data))
+
+    @json_app.route("/")
+    @paginate(limit=1, skip=0, headers=True)
+    def foo():
+        return Cursor([1, 2, 3])
+
+    response = client.get("/")
+    assert response.status_code == 200
+
+    assert "X-Total-Count" in response.headers
+    assert "Link" in response.headers
+
+
 def test_stupeflask_converters(json_app, client):
     @json_app.route("/<ObjectId:foo_id>")
     def foo_id(foo_id):

tests/pagination.py

@@ -75,3 +75,30 @@ def test_paginate_function(app):
         def foo_instance():
             return Cursor([1, 2, 3])
         assert paginate(foo_instance, skip=2)().data == [3]
+
+
+@pytest.mark.parametrize("app", [
+    Stupeflask(__name__),
+    Flask(__name__)
+])
+def test_paginate_header_total_count(app):
+    with app.test_request_context():
+        paginate(Cursor([1, 2, 3]), headers=True)
+        if isinstance(app, Stupeflask):
+            assert request.response_headers["X-Total-Count"] == 3
+            assert "Link" not in request.response_headers
+
+
+@pytest.mark.parametrize("app", [
+    Stupeflask(__name__),
+    Flask(__name__)
+])
+def test_paginate_header_link(app):
+    with app.test_request_context():
+        paginate(Cursor([1, 2, 3]), limit=1, skip=0, headers=True)
+        if isinstance(app, Stupeflask):
+            links = request.response_headers["Link"].split(",")
+            assert "first" in links[0]
+            assert "prev" in links[1]
+            assert "next" in links[2]
+            assert "last" in links[3]