WIP: Implement pagination headers

Fixes #19

TODO:
- Implement & test header Link for bare Flask app

flask_stupe/app.py

@@ -1,7 +1,7 @@
 import os
 from pkgutil import iter_modules
 
-from flask import Blueprint, Flask
+from flask import Blueprint, Flask, Response, request
 
 from flask_stupe.config import Config
 from flask_stupe.converters import converters
@@ -24,6 +24,14 @@ def __init__(self, *args, **kwargs):
         log.info(" * Overriden by environment: " + ", ".join(from_env))
         self.register_converters(converters)
 
+    def make_response(self, rv):
+        rv = Response(rv)
+
+        if request.response_headers:
+            rv.headers.extend(request.response_headers)
+
+        return rv
+
     def register_converter(self, converter, name=None):
         """Register a new converter that can be used in endpoints URLs
 

flask_stupe/json.py

@@ -124,6 +124,7 @@ def make_response(self, rv):
 
         rv = jsonify(rv)
         rv.status_code = code
+
         return rv
 
     def __init__(self, *args, **kwargs):

flask_stupe/pagination.py

@@ -1,4 +1,5 @@
 import functools
+from collections import OrderedDict
 
 from flask import request
 
@@ -8,10 +9,47 @@
 
 
 if pymongo:
+
+    def _get_pagination_links(skip, limit, total_count=None):
+        template = "{}?limit={}&skip={{skip}}".format(request.base_url, limit)
+        links = OrderedDict([
+            ("self", template.format(skip=skip)),
+            ("first", template.format(skip=0))
+        ])
+        prev_skip = skip - limit
+        if prev_skip >= 0:
+            links.update(prev=template.format(skip=prev_skip))
+        next_skip = skip + limit
+        if next_skip < total_count:
+            links.update(next=template.format(skip=next_skip))
+        if total_count:
+            links.update(last=template.format(skip=total_count - limit))
+        return links
+
     def _paginate(cursor, skip=None, limit=None, sort=None, count=True):
+        total_count = None
+        if count:
+            total_count = cursor.count()
+        links = None
+        if limit:
+            links = _get_pagination_links(skip or 0, limit, total_count)
+
+        headers = getattr(request, "response_headers", None)
+        if isinstance(headers, dict):
+            if total_count:
+                headers["X-Total-Count"] = total_count
+            if links:
+                header_links = []
+                for name, link in links.items():
+                    header_links.append('<{}>; rel="{}"'.format(link, name))
+                headers["Link"] = ", ".join(header_links)
+
         metadata = getattr(request, "metadata", None)
-        if count and isinstance(metadata, dict):
-            metadata.update(count=cursor.count())
+        if isinstance(metadata, dict):
+            if total_count:
+                metadata.update(count=total_count)
+            if links:
+                metadata.update(links=links)
 
         skip = request.args.get("skip", skip, type=int)
         if skip is not None:

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/app.py

@@ -5,6 +5,8 @@
 
 from flask_stupe.app import Stupeflask
 from flask_stupe.config import Config
+from flask_stupe.pagination import paginate
+from tests.conftest import Cursor
 
 
 def test_app_config(app):
@@ -53,3 +55,19 @@ def test_register_blueprints(test_apps, caplog, app):
 @pytest.fixture
 def app():
     return Stupeflask(__name__)
+
+
+def test_stupeflask_response_with_paginate_headers(app, client):
+
+    @app.route("/")
+    def foo():
+        cursor = Cursor([1, 2, 3])
+        cursor = paginate(cursor, limit=1, skip=0)
+        return (lambda c: c.data)(cursor)
+
+    response = client.get("/")
+    print(response.data)
+    assert response.status_code == 200
+
+    assert "X-Total-Count" in response.headers
+    assert "Link" in response.headers

tests/pagination.py

@@ -75,3 +75,32 @@ def test_paginate_function(app):
         def foo_instance():
             return Cursor([1, 2, 3])
         assert paginate(foo_instance, skip=2)().data == [3]
+
+
+def test_paginate_header_total_count(app):
+    with app.test_request_context():
+        paginate(Cursor([1, 2, 3]))
+        assert request.response_headers["X-Total-Count"] == 3
+        assert "Link" not in request.response_headers
+
+
+def test_paginate_header_link(app):
+    with app.test_request_context():
+        paginate(Cursor([1, 2, 3]), limit=1, skip=1)
+        links = request.response_headers["Link"].split(",")
+        assert links[0].split("?")[1] == 'limit=1&skip=1>; rel="self"'
+        assert links[1].split("?")[1] == 'limit=1&skip=0>; rel="first"'
+        assert links[2].split("?")[1] == 'limit=1&skip=0>; rel="prev"'
+        assert links[3].split("?")[1] == 'limit=1&skip=2>; rel="next"'
+        assert links[4].split("?")[1] == 'limit=1&skip=2>; rel="last"'
+
+
+def test_paginate_metadata_links(app):
+    with app.test_request_context():
+        paginate(Cursor([1, 2, 3]), limit=1, skip=1)
+        links = request.metadata["links"]
+        assert links["self"].split("?")[1] == "limit=1&skip=1"
+        assert links["first"].split("?")[1] == "limit=1&skip=0"
+        assert links["prev"].split("?")[1] == "limit=1&skip=0"
+        assert links["next"].split("?")[1] == "limit=1&skip=2"
+        assert links["last"].split("?")[1] == "limit=1&skip=2"