From 7526b7bdef042b5cf3bedb18d709ec6e8ba0d0d7 Mon Sep 17 00:00:00 2001
From: Roman Masse <roman@1000mercis.com>
Date: Thu, 30 Aug 2018 12:03:54 +0200
Subject: [PATCH] Implement pagination headers

Fixes #19
---
 flask_stupe/json.py       |  4 ++++
 flask_stupe/pagination.py | 34 ++++++++++++++++++++++++++--------
 flask_stupe/request.py    |  1 +
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/flask_stupe/json.py b/flask_stupe/json.py
index 22f0661..e86293d 100644
--- a/flask_stupe/json.py
+++ b/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):
diff --git a/flask_stupe/pagination.py b/flask_stupe/pagination.py
index 658ad56..6adcffb 100644
--- a/flask_stupe/pagination.py
+++ b/flask_stupe/pagination.py
@@ -1,17 +1,34 @@
 import functools
 
 from flask import request
-
 from flask_stupe import pymongo
 
 __all__ = []
 
 
 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(total_count, limit, skip):
+        template = '<{}?limit={}&skip={{skip}}>; ' \
+                   'rel="{{rel}}"'.format(request.base_url, limit)
+        return ", ".join([
+            template.format(skip=0, rel="first"),
+            template.format(skip=skip - limit, rel="prev"),
+            template.format(skip=skip + limit, rel="next"),
+            template.format(skip=total_count - limit, rel="last")
+        ])
+
+    def _paginate(cursor, skip=None, limit=None, sort=None, count=True,
+                  headers=False):
+        if headers:
+            total_count = cursor.count()
+            headers = getattr(request, "response_headers")
+            headers["X-Total-Count"] = total_count
+            headers["Link"] = _get_link_header(total_count, limit, skip)
+        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 +52,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:
diff --git a/flask_stupe/request.py b/flask_stupe/request.py
index de1fbd6..f88cc66 100644
--- a/flask_stupe/request.py
+++ b/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"]