Repository: apiflask/apiflask (python)

Issue:
Add support for adding decorators to the OpenAPI spec and documentation UI endpoints
As a user of the APIFlask framework, I would like the ability to apply a decorator to the endpoints used to serve the OpenAPI spec as well as the documentation UI, so that I can execute additional logic before returning back the OpenAPI spec or the documentation UI to a client.

The more specific use case I have for this is to apply some security checks to validate that the client trying to access the documentation is allowed to access the endpoints that serve it.

I have currently achieved this by doing the following, but it's rather ugly - I had to look into the code in  `APIFlask._register_openapi_blueprint` and replicate some of its logic to rewrite the view functions, which Flask advises against doing:

```
def _add_authentication_to_docs(app: APIFlask) -> None:
    if app.spec_path:
        _rewrite_view_function(app, "spec")
    if app.docs_path:
        _rewrite_view_function(app, "docs")
        if app.docs_ui == "swagger-ui" and app.docs_oauth2_redirect_path:
            _rewrite_view_function(app, "swagger_ui_oauth_redirect")


def _rewrite_view_function(app: APIFlask, method_name: str) -> None:
    # Flask warns you not to do this, but we don't really have a choice...
    app.view_functions[f"openapi.{method_name}"] = authenticate_apidocs()(app.view_functions[f"openapi.{method_name}"])
```

APIFairy provides this functionality via the `APIFAIRY_APISPEC_DECORATORS` and `APIFAIRY_UI_DECORATORS` config variables, which is much more straightforward:

```
app.config["APIFAIRY_APISPEC_DECORATORS"] = [authenticate_apidocs()]
app.config["APIFAIRY_UI_DECORATORS"] = [authenticate_apidocs()]
```

I am hoping that this framework might consider adopting similar functionality.

Pull request description:
[Feature] support decorators to the openapi
fixes #483 

<!--
Link to relevant issues or previous PRs, one per line. Use "fixes" to
automatically close an issue.
-->



Checklist:

- [x] Add tests that demonstrate the correct behavior of the change. Tests should fail without the change.
- [x] Add or update relevant docs, in the `docs` folder and in code docstring.
- [x] Add an entry in `CHANGES.md` summarizing the change and linking to the issue.
- [x] Add `*Version changed*` or `*Version added*` note in any relevant docs and docstring.
- [x] Run `pytest` and `tox`, no tests failed.


Patch:
diff --git a/CHANGES.md b/CHANGES.md
index 7c0d550..dd395d8 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -3,9 +3,16 @@
 Released: -
 
 - Allow adding multiple media types for a response ([issue #494][issue_494])
+- Support adding headers to the response schema ([issue #507][issue_507])
 
 [issue_494]: https://github.com/apiflask/apiflask/issues/494
+[issue_507]: https://github.com/apiflask/apiflask/issues/507
 
+## Version 2.0.3
+
+- Add support for adding decorators to the OpenAPI spec and documentation UI endpoints ([issue #483][issue_483]).
+
+[issue_483]: https://github.com/apiflask/apiflask/issues/483
 
 ## Version 2.0.2
 
diff --git a/docs/api-docs.md b/docs/api-docs.md
index c78a88f..222d0af 100644
--- a/docs/api-docs.md
+++ b/docs/api-docs.md
@@ -64,6 +64,11 @@ Now the paths to docs and spec will be <http://localhost:5000/openapi/docs>
 and <http://localhost:5000/openapi/openapi.json>.
 
 
+## Add authentication to API documentation
+
+See *[Protect OpenAPI endpoints](/openapi/#protect-openapi-endpoints)* for more details.
+
+
 ## Add custom API documentation
 
 You can easily add support to other API docs or serve the supported docs UI by yourself.
diff --git a/docs/configuration.md b/docs/configuration.md
index 023e550..892f97e 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -577,6 +577,45 @@ app.config['SPEC_PROCESSOR_PASS_OBJECT'] = True
     This configuration variable was added in the [version 1.3.0](/changelog/#version-130).
 
 
+### DOCS_DECORATORS
+
+The custom decorators of the OpenAPI documentation UI endpoint (`/docs`).
+
+- Type: `List[Callable]`
+- Default value: `None`
+- Examples:
+
+```python
+app.config['DOCS_DECORATORS'] = [your_decorator]
+```
+
+
+### SPEC_DECORATORS
+
+The custom decorators of the OpenAPI spec endpoint (`/openapi.json`).
+
+- Type: `List[Callable]`
+- Default value: `None`
+- Examples:
+
+```python
+app.config['SPEC_DECORATORS'] = [your_decorator]
+```
+
+
+### SWAGGER_UI_OAUTH_REDIRECT_DECORATORS
+
+The custom decorators of the Swagger UI OAuth redirect endpoint (`/docs/oauth2-redirect`).
+
+- Type: `List[Callable]`
+- Default value: `None`
+- Examples:
+
+```python
+app.config['SWAGGER_UI_OAUTH_REDIRECT_DECORATORS'] = [your_decorator]
+```
+
+
 ## Automation behavior control
 
 The following configuration variables are used to control the automation behavior
diff --git a/docs/openapi.md b/docs/openapi.md
index 2d920c2..67aa6bc 100644
--- a/docs/openapi.md
+++ b/docs/openapi.md
@@ -715,6 +715,36 @@ def new_pet():
 ```
 
 
+## Response `headers`
+
+!!! warning "Version >= 2.0.3"
+
+    This feature was added in the [version 2.0.3](/changelog/#version-203).
+
+You can pass a schema class to `headers` in the output decorator:
+```python
+class Header(Schema):
+    x_token = String(
+        data_key='X-Token',
+        metadata={'description': 'A custom token header'}
+    )
+
+@app.post('/pets')
+@app.output(PetOut, headers=Header)
+def new_pet():
+    pass
+```
+
+When passing the schema to `headers`, you can also use a dict instead of a schema class:
+
+```python
+@app.post('/pets')
+@app.output(PetOut, headers={'X-Token': String(metadata={'description': 'A custom token header'})})
+def new_pet():
+    pass
+```
+
+
 ## Request and response content type / media type
 
 For request, the content type is set automatically based on the input location:
@@ -989,6 +1019,20 @@ auth = HTTPTokenAuth(description='some description')
 ```
 
 
+## Protect OpenAPI endpoints
+
+If you want to apply an authentication or custom handling logic
+to the OpenAPI endpoints, you can use the `SPEC_DECORATORS`,
+`DOCS_DECORATORS` and `SWAGGER_UI_OAUTH_REDIRECT_DECORATORS` configuration options:
+
+```python
+app.config['SPEC_DECORATORS'] = [app.auth_required(auth)]
+app.config['DOCS_DECORATORS'] = [app.auth_required(auth)]
+```
+
+Check out a full example at [examples/openapi/custom_decorators/app.py](https://github.com/apiflask/apiflask/tree/main/examples/openapi/custom_decorators/app.py)
+
+
 ## Disable the OpenAPI support
 
 
diff --git a/examples/README.md b/examples/README.md
index 753ce31..1c57b65 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -63,7 +63,8 @@ Each example application store in a sub-folder:
 - `/cbv`: Classed-based view example
 - `/orm`: ORM example (with Flask-SQLAlchemy)
 - `/pagination`: Pagination example (with Flask-SQLAlchemy)
-- `/openapi`: OpenAPI example
+- `/openapi/basic`: Basic OpenAPI example
+- `/openapi/custom_decorators`: OpenAPI example with custom decorators
 - `/base_response`: Base response example
 - `/dataclass`: Dataclass example (with marshmallow-dataclass)
 
diff --git a/examples/dataclass/app.py b/examples/dataclass/app.py
index 527785e..428f647 100644
--- a/examples/dataclass/app.py
+++ b/examples/dataclass/app.py
@@ -10,10 +10,21 @@ app = APIFlask(__name__)
 @dataclass
 class PetIn:
     name: str = field(
-        metadata={'required': True, 'validate': Length(min=1, max=10)}
+        metadata={
+            'required': True,
+            'validate': Length(min=1, max=10),
+            'example': 'Medor',
+            'description': 'This will be printed in the generated doc. '
+                           'The "example" value "Medor" will be fed '
+                           'into the "try it"/"Send API request".',
+        }
     )
     category: str = field(
-        metadata={'required': True, 'validate': OneOf(['cat', 'dog'])}
+        default='dog',
+        metadata={
+            'required': True,
+            'validate': OneOf(['cat', 'dog'])
+        }
     )
 
 
diff --git a/examples/openapi/app.py b/examples/openapi/basic/app.py
similarity index 100%
rename from examples/openapi/app.py
rename to examples/openapi/basic/app.py
diff --git a/examples/openapi/custom_decorators/app.py b/examples/openapi/custom_decorators/app.py
new file mode 100644
index 0000000..6f6749b
--- /dev/null
+++ b/examples/openapi/custom_decorators/app.py
@@ -0,0 +1,39 @@
+import typing as t
+
+from apiflask import APIFlask, HTTPBasicAuth
+from werkzeug.security import generate_password_hash, check_password_hash
+
+app = APIFlask(__name__)
+auth = HTTPBasicAuth()
+
+users = {
+    'userA': generate_password_hash('foo'),
+    'userB': generate_password_hash('bar'),
+}
+
+
+# if you want to apply an authentication or other decorator
+# to the endpoints used to serve the OpenAPI spec
+# as well as the documentation UI, use the
+# `SPEC_DECORATORS`, `DOCS_DECORATORS` and
+# `SWAGGER_UI_OAUTH_REDIRECT_DECORATORS` configuration options
+app.config['SPEC_DECORATORS'] = [app.auth_required(auth)]
+app.config['DOCS_DECORATORS'] = [app.auth_required(auth)]
+
+
+@auth.verify_password
+def verify_password(username: str, password: str) -> t.Union[str, None]:
+    if (
+        username in users
+        and check_password_hash(users[username], password)
+    ):
+        return username
+    return None
+
+
+@app.route('/')
+def index():
+    return {
+        'spec': '/openapi.json',
+        'docs': '/docs',
+    }
diff --git a/src/apiflask/app.py b/src/apiflask/app.py
index 18ac387..038ad4f 100644
--- a/src/apiflask/app.py
+++ b/src/apiflask/app.py
@@ -3,6 +3,7 @@ import json
 import re
 import typing as t
 import warnings
+from functools import wraps
 
 from apispec import APISpec
 from apispec import BasePlugin
@@ -522,6 +523,10 @@ class APIFlask(APIScaffold, Flask):
         The name of the blueprint is "openapi". This blueprint will hold the view
         functions for spec file and API docs.
 
+        *Version changed: 2.0.3*
+
+        - Inject the OpenAPI endpoints decorators.
+
         *Version changed: 1.1.0*
 
         - Deprecate the redoc view at /redoc path.
@@ -538,6 +543,7 @@ class APIFlask(APIScaffold, Flask):
 
         if self.spec_path:
             @bp.route(self.spec_path)
+            @self._apply_decorators(config_name='SPEC_DECORATORS')
             def spec():
                 if self.config['SPEC_FORMAT'] == 'json':
                     response = jsonify(self._get_spec('json'))
@@ -555,6 +561,7 @@ class APIFlask(APIScaffold, Flask):
                 )
 
             @bp.route(self.docs_path)
+            @self._apply_decorators(config_name='DOCS_DECORATORS')
             def docs():
                 return render_template_string(
                     ui_templates[self.docs_ui],
@@ -566,6 +573,7 @@ class APIFlask(APIScaffold, Flask):
             if self.docs_ui == 'swagger-ui':
                 if self.docs_oauth2_redirect_path:
                     @bp.route(self.docs_oauth2_redirect_path)
+                    @self._apply_decorators(config_name='SWAGGER_UI_OAUTH_REDIRECT_DECORATORS')
                     def swagger_ui_oauth_redirect() -> str:
                         return render_template_string(swagger_ui_oauth2_redirect_template)
 
@@ -1017,6 +1025,7 @@ class APIFlask(APIScaffold, Flask):
                     examples = view_func._spec.get('response')['examples']
                     links = view_func._spec.get('response')['links']
                     content_type = view_func._spec.get('response')['content_type']
+                    headers = view_func._spec.get('response')['headers']
                     add_response(
                         operation,
                         status_code,
@@ -1026,6 +1035,8 @@ class APIFlask(APIScaffold, Flask):
                         examples,
                         links,
                         content_type,
+                        headers,
+                        ma_plugin,
                     )
                 else:
                     # add a default 200 response for views without using @app.output
@@ -1193,3 +1204,21 @@ class APIFlask(APIScaffold, Flask):
             spec.path(path=path, operations=sorted_operations)
 
         return spec
+
+    def _apply_decorators(self, config_name: str):
+        """Apply the decorators to the OpenAPI endpoints at runtime.
+
+        Arguments:
+            config_name: The config name to get the list of decorators.
+        """
+        def decorator(f):
+            @wraps(f)
+            def wrapper(*args, **kwargs):
+                decorated_func = f
+                decorators = self.config[config_name]
+                if decorators:
+                    for decorator in decorators:
+                        decorated_func = decorator(decorated_func)
+                return decorated_func(*args, **kwargs)
+            return wrapper
+        return decorator
diff --git a/src/apiflask/openapi.py b/src/apiflask/openapi.py
index eec1a71..3c90642 100644
--- a/src/apiflask/openapi.py
+++ b/src/apiflask/openapi.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import typing as t
 
 from apispec import APISpec
+from apispec.ext.marshmallow import MarshmallowPlugin
 
 from .exceptions import _bad_schema_message
 from .schemas import EmptySchema
@@ -34,6 +35,7 @@ default_response = {
     'examples': None,
     'links': None,
     'content_type': 'application/json',
+    'headers': None,
 }
 
 
@@ -167,9 +169,15 @@ def add_response(
     examples: t.Optional[t.Dict[str, t.Any]] = None,
     links: t.Optional[t.Dict[str, t.Any]] = None,
     content_type: t.Optional[str] = 'application/json',
+    headers_schema: t.Optional[SchemaType] = None,
+    ma_plugin: t.Optional[MarshmallowPlugin] = None,
 ) -> None:
     """Add response to operation.
 
+    *Version changed: 2.0.3*
+
+    - Add parameter `headers_schema`.
+
     *Version changed: 1.3.0*
 
     - Add parameter `content_type`.
@@ -198,6 +206,14 @@ def add_response(
             content_type]['examples'] = examples
     if links is not None:
         operation['responses'][status_code]['links'] = links
+    if headers_schema is not None:
+        headers = ma_plugin.converter.schema2parameters(  # type: ignore
+            headers_schema,
+            location='headers'
+        )
+        operation['responses'][status_code]['headers'] = {
+            header['name']: header for header in headers
+        }
 
 
 def add_response_with_schema(
diff --git a/src/apiflask/scaffold.py b/src/apiflask/scaffold.py
index 2f12609..d6ce574 100644
--- a/src/apiflask/scaffold.py
+++ b/src/apiflask/scaffold.py
@@ -338,6 +338,7 @@ class APIScaffold:
         examples: t.Optional[t.Dict[str, t.Any]] = None,
         links: t.Optional[t.Dict[str, t.Any]] = None,
         content_type: t.Optional[str] = 'application/json',
+        headers: t.Optional[SchemaType] = None,
     ) -> t.Callable[[DecoratedType], DecoratedType]:
         """Add output settings for view functions.
 
@@ -406,6 +407,11 @@ class APIScaffold:
                 about setting response links.
 
             content_type: The content/media type of the response. It defautls to `application/json`.
+            headers: The schemas of the headers.
+
+        *Version changed: 2.0.3*
+
+        - Add parameter `headers`.
 
         *Version changed: 2.0.0*
 
@@ -446,6 +452,14 @@ class APIScaffold:
         if isinstance(schema, type):  # pragma: no cover
             schema = schema()
 
+        if headers is not None:
+            if headers == {}:
+                headers = EmptySchema
+            if isinstance(headers, ABCMapping):
+                headers = _generate_schema_from_mapping(headers, None)
+            if isinstance(headers, type):
+                headers = headers()
+
         def decorator(f):
             f = _ensure_sync(f)
             _annotate(f, response={
@@ -456,6 +470,7 @@ class APIScaffold:
                 'examples': examples,
                 'links': links,
                 'content_type': content_type,
+                'headers': headers,
             })
 
             def _jsonify(
diff --git a/src/apiflask/settings.py b/src/apiflask/settings.py
index 7951a3f..9f29bcd 100644
--- a/src/apiflask/settings.py
+++ b/src/apiflask/settings.py
@@ -25,6 +25,9 @@ LOCAL_SPEC_PATH: t.Optional[str] = None
 LOCAL_SPEC_JSON_INDENT: int = 2
 SYNC_LOCAL_SPEC: t.Optional[bool] = None
 SPEC_PROCESSOR_PASS_OBJECT: bool = False
+SPEC_DECORATORS: t.Optional[t.List[t.Callable]] = None
+DOCS_DECORATORS: t.Optional[t.List[t.Callable]] = None
+SWAGGER_UI_OAUTH_REDIRECT_DECORATORS: t.Optional[t.List[t.Callable]] = None
 # Automation behavior control
 AUTO_TAGS: bool = True
 AUTO_SERVERS: bool = True


Test patch:
diff --git a/examples/test_examples.py b/examples/test_examples.py
index 254a4f8..4c350a1 100644
--- a/examples/test_examples.py
+++ b/examples/test_examples.py
@@ -9,7 +9,7 @@ full_examples = [
     'basic',
     'orm',
     'cbv',
-    'openapi',
+    'openapi/basic',
     'blueprint_tags',
     'base_response',
     'dataclass',
diff --git a/tests/schemas.py b/tests/schemas.py
index eec2a33..b58d806 100644
--- a/tests/schemas.py
+++ b/tests/schemas.py
@@ -86,3 +86,11 @@ class CustomHTTPError(Schema):
     status_code = String(required=True)
     message = String(required=True)
     custom = String(required=True)
+
+
+class ResponseHeader(Schema):
+    x_token = String(
+        data_key='X-Token',
+        required=True,
+        metadata={'description': 'A custom token header'}
+    )
diff --git a/tests/test_app.py b/tests/test_app.py
index 9ae6406..06c5a04 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -10,6 +10,7 @@ from .schemas import Pagination
 from apiflask import APIBlueprint
 from apiflask import APIFlask
 from apiflask import Schema
+from apiflask.exceptions import abort
 from apiflask.fields import Integer
 from apiflask.fields import String
 
@@ -313,3 +314,54 @@ def test_apispec_plugins(app):
     spec = app._get_spec('json')
 
     assert spec['paths']['/plugin_test'].get('post') == 'some_injected_test_data'
+
+
+def test_spec_decorators(app, client):
+    def auth_decorator(f):
+        def wrapper(*args, **kwargs):
+            abort(401)
+            return f(*args, **kwargs)
+        return wrapper
+
+    app.config['SPEC_DECORATORS'] = [auth_decorator]
+
+    rv = client.get('/openapi.json')
+    assert rv.status_code == 401
+
+    rv = client.get('/docs')
+    assert rv.status_code == 200
+
+
+def test_docs_decorators(app, client):
+    def auth_decorator(f):
+        def wrapper(*args, **kwargs):
+            abort(401)
+            return f(*args, **kwargs)
+        return wrapper
+
+    app.config['DOCS_DECORATORS'] = [auth_decorator]
+
+    rv = client.get('/openapi.json')
+    assert rv.status_code == 200
+
+    rv = client.get('/docs')
+    assert rv.status_code == 401
+
+
+def test_swagger_ui_oauth_redirect_decorators(app, client):
+    def auth_decorator(f):
+        def wrapper(*args, **kwargs):
+            abort(401)
+            return f(*args, **kwargs)
+        return wrapper
+
+    app.config['SWAGGER_UI_OAUTH_REDIRECT_DECORATORS'] = [auth_decorator]
+
+    rv = client.get('/openapi.json')
+    assert rv.status_code == 200
+
+    rv = client.get('/docs')
+    assert rv.status_code == 200
+
+    rv = client.get('/docs/oauth2-redirect')
+    assert rv.status_code == 401
diff --git a/tests/test_openapi_headers.py b/tests/test_openapi_headers.py
new file mode 100644
index 0000000..5355f58
--- /dev/null
+++ b/tests/test_openapi_headers.py
@@ -0,0 +1,89 @@
+from .schemas import Foo
+from .schemas import ResponseHeader
+from apiflask.fields import Boolean
+from apiflask.fields import Integer
+from apiflask.fields import List
+from apiflask.fields import Number
+from apiflask.fields import String
+
+
+def test_spec_with_dict_headers(app, client):
+    @app.route('/foo')
+    @app.output(Foo, headers={
+        'X-boolean': Boolean(metadata={'description': 'A boolean header'}),
+        'X-integer': Integer(metadata={'description': 'An integer header'}),
+        'X-number': Number(metadata={'description': 'A number header'}),
+        'X-string': String(metadata={'description': 'A string header'}),
+        'X-array': List(String(), metadata={'description': 'An array header'}),
+    })
+    def foo():
+        pass
+
+    rv = client.get('/openapi.json')
+    assert rv.json['paths']['/foo']['get']['responses']['200']['headers'] == {
+        'X-boolean': {
+            'name': 'X-boolean',
+            'in': 'header',
+            'description': 'A boolean header',
+            'required': False,
+            'schema': {'type': 'boolean'},
+        },
+        'X-integer': {
+            'name': 'X-integer',
+            'in': 'header',
+            'description': 'An integer header',
+            'required': False,
+            'schema': {'type': 'integer'},
+        },
+        'X-number': {
+            'name': 'X-number',
+            'in': 'header',
+            'description': 'A number header',
+            'required': False,
+            'schema': {'type': 'number'},
+        },
+        'X-string': {
+            'name': 'X-string',
+            'in': 'header',
+            'description': 'A string header',
+            'required': False,
+            'schema': {'type': 'string'},
+        },
+        'X-array': {
+            'name': 'X-array',
+            'in': 'header',
+            'description': 'An array header',
+            'required': False,
+            'schema': {'items': {'type': 'string'}, 'type': 'array'},
+            'style': 'form',
+            'explode': True
+        },
+    }
+
+
+def test_spec_with_empty_headers(app, client):
+    @app.route('/foo')
+    @app.output(Foo, headers={})
+    def foo():
+        pass
+
+    rv = client.get('/openapi.json')
+    assert rv.json['paths']['/foo']['get']['responses']['200']['headers'] == {}
+
+
+def test_spec_with_schema_headers(app, client):
+    @app.route('/foo')
+    @app.output(Foo, headers=ResponseHeader)
+    def foo():
+        pass
+
+    rv = client.get('/openapi.json')
+    assert rv.json['paths']['/foo']['get']['responses']['200']['headers'] == {
+        'X-Token': {
+            'name': 'X-Token',
+            'in': 'header',
+            'description': 'A custom token header',
+            'required': True,
+            'schema': {'type': 'string'},
+        },
+    }