Document pagination header in spec

marshmallow-code · Mar 29, 2019 · 134f9c3 · 134f9c3
1 parent 8e86516
commit 134f9c3
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 0 deletions.
diff --git a/flask_rest_api/pagination.py b/flask_rest_api/pagination.py
@@ -133,6 +133,11 @@ class PaginationMixin:
     DEFAULT_PAGINATION_PARAMETERS = {
         'page': 1, 'page_size': 10, 'max_page_size': 100}
 
+    PAGINATION_HEADER_DOC = {
+        'description': 'Pagination metadata',
+        'schema': PaginationHeaderSchema,
+    }
+
     def paginate(self, pager=None, *,
                  page=None, page_size=None, max_page_size=None):
         """Decorator adding pagination to the endpoint
@@ -171,6 +176,7 @@ def decorator(func):
             # Add pagination params to doc info in function object
             func._apidoc = getattr(func, '_apidoc', {})
             func._apidoc.setdefault('parameters', []).append(parameters)
+            func._paginated = True
 
             @wraps(func)
             def wrapper(*args, **kwargs):

diff --git a/flask_rest_api/response.py b/flask_rest_api/response.py
@@ -44,6 +44,13 @@ def decorator(func):
             doc_schema = self._make_doc_response_schema(schema)
             if doc_schema:
                 doc['responses'][str(code)]['schema'] = doc_schema
+            # Document pagination header if needed
+            if getattr(func, '_paginated', False) is True:
+                doc['responses'][str(code)]['headers'] = {
+                    self.PAGINATION_HEADER_FIELD_NAME: (
+                        self.PAGINATION_HEADER_DOC
+                    )
+                }
             func._apidoc = deepupdate(getattr(func, '_apidoc', {}), doc)
 
             @wraps(func)

diff --git a/tests/test_pagination.py b/tests/test_pagination.py
@@ -150,6 +150,31 @@ def func(pagination_parameters):
                 '"first_page": 1, "last_page": 1, "page": 1}'
             )
 
+    def test_pagination_header_documentation(self, app):
+        """Test pagination header is documented"""
+        api = Api(app)
+
+        class CustomBlueprint(Blueprint):
+            PAGINATION_HEADER_FIELD_NAME = 'X-Custom-Pagination-Header'
+
+        blp = CustomBlueprint('test', __name__, url_prefix='/test')
+
+        @blp.route('/')
+        @blp.response()
+        @blp.paginate()
+        def func(pagination_parameters):
+            """Dummy view func"""
+
+        api.register_blueprint(blp)
+        spec = api.spec.to_dict()
+        get = spec['paths']['/test/']['get']
+        assert get['responses']['200']['headers'] == {
+            'X-Custom-Pagination-Header': {
+                'description': 'Pagination metadata',
+                'schema': {'$ref': '#/components/schemas/PaginationHeader'},
+            }
+        }
+
     @pytest.mark.parametrize('header_name', ('X-Pagination', None))
     def test_pagination_item_count_missing(self, app, header_name):
         """If item_count was not set, pass and warn"""