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..2b2bd1d 100644 --- a/flask_stupe/pagination.py +++ b/flask_stupe/pagination.py @@ -1,17 +1,39 @@ 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(base_url, total_count, limit, skip): + links = [] + template = '<{}?limit={{limit}}&skip={{skip}}>; ' \ + 'rel="{{rel}}"'.format(base_url) + prev_skip = skip - limit + next_skip = skip + limit + if prev_skip >= 0: + k = "prev" if (prev_skip > limit) else "first" + links.append(template.format(limit=limit, skip=skip, rel=k)) + if next_skip < total_count: + k = "next" if (next_skip < (total_count - limit)) else "last" + links.append(template.format(limit=limit, skip=skip, rel=k)) + return ", ".join(links) + + 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(request.base_url, 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 +57,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"]