Merge pull request #59 from hkwi/flask_blueprint

Add flask Blueprint support

spectree/plugins/flask_plugin.py

@@ -5,21 +5,38 @@
 
 
 class FlaskPlugin(BasePlugin):
+    blueprint_state = None
+
     def find_routes(self):
-        for rule in self.app.url_map.iter_rules():
-            if any(str(rule).startswith(path) for path in (
-                    f'/{self.config.PATH}', '/static'
-            )):
-                continue
-            yield rule
+        from flask import current_app
+        if self.blueprint_state:
+            excludes = [f"{self.blueprint_state.blueprint.name}.{ep}"
+                for ep in ["static", "openapi"]+[f"doc_page_{ui}" for ui in PAGES]]
+            for rule in current_app.url_map.iter_rules():
+                if self.blueprint_state.url_prefix and not str(rule).startswith(self.blueprint_state.url_prefix):
+                    continue
+                if rule.endpoint in excludes:
+                    continue
+                yield rule
+        else:
+            for rule in self.app.url_map.iter_rules():
+                if any(str(rule).startswith(path) for path in (
+                        f'/{self.config.PATH}', '/static'
+                )):
+                    continue
+                yield rule
 
     def bypass(self, func, method):
         if method in ['HEAD', 'OPTIONS']:
             return True
         return False
 
     def parse_func(self, route):
-        func = self.app.view_functions[route.endpoint]
+        if self.blueprint_state:
+            func = self.blueprint_state.app.view_functions[route.endpoint]
+        else:
+            func = self.app.view_functions[route.endpoint]
+
         for method in route.methods:
             yield method, func
 
@@ -141,17 +158,36 @@ def validate(self,
 
     def register_route(self, app):
         self.app = app
-        from flask import jsonify
+        from flask import jsonify, Blueprint
 
         self.app.add_url_rule(
             self.config.spec_url,
             'openapi',
             lambda: jsonify(self.spectree.spec),
         )
 
-        for ui in PAGES:
-            self.app.add_url_rule(
-                f'/{self.config.PATH}/{ui}',
-                f'doc_page_{ui}',
-                lambda ui=ui: PAGES[ui].format(self.config.spec_url)
-            )
+        if isinstance(app, Blueprint):
+            def gen_doc_page(ui):
+                state = self.blueprint_state
+
+                spec_url = self.config.spec_url
+                if state.url_prefix is not None:
+                    spec_url = "/".join((state.url_prefix.rstrip("/"), self.config.spec_url.lstrip("/")))
+
+                return PAGES[ui].format(spec_url)
+
+            for ui in PAGES:
+                app.add_url_rule(
+                    f'/{self.config.PATH}/{ui}',
+                    f'doc_page_{ui}',
+                    lambda ui=ui: gen_doc_page(ui)
+                )
+
+            app.record(lambda state: setattr(self, "blueprint_state", state))
+        else:
+            for ui in PAGES:
+                self.app.add_url_rule(
+                    f'/{self.config.PATH}/{ui}',
+                    f'doc_page_{ui}',
+                    lambda ui=ui: PAGES[ui].format(self.config.spec_url)
+                )

tests/test_plugin_flask_blueprint.py

@@ -0,0 +1,105 @@
+from random import randint
+import pytest
+import json
+from flask import Flask, Blueprint, jsonify, request
+
+from spectree import SpecTree, Response
+
+from .common import Query, Resp, JSON, Headers, Cookies
+
+
+def before_handler(req, resp, err, _):
+    if err:
+        resp.headers['X-Error'] = 'Validation Error'
+
+
+def after_handler(req, resp, err, _):
+    resp.headers['X-Validation'] = 'Pass'
+
+
+def api_after_handler(req, resp, err, _):
+    resp.headers['X-API'] = 'OK'
+
+
+api = SpecTree('flask', before=before_handler, after=after_handler)
+app = Blueprint('test_blueprint', __name__)
+
+
+@app.route('/ping')
+@api.validate(headers=Headers, tags=['test', 'health'])
+def ping():
+    """summary
+    description"""
+    return jsonify(msg='pong')
+
+
+@app.route('/api/user/<name>', methods=['POST'])
+@api.validate(
+    query=Query,
+    json=JSON,
+    cookies=Cookies,
+    resp=Response(HTTP_200=Resp, HTTP_401=None),
+    tags=['api', 'test'],
+    after=api_after_handler)
+def user_score(name):
+    score = [randint(0, request.context.json.limit) for _ in range(5)]
+    score.sort(reverse=request.context.query.order)
+    assert request.context.cookies.pub == 'abcdefg'
+    assert request.cookies['pub'] == 'abcdefg'
+    return jsonify(name=request.context.json.name, score=score)
+
+
+api.register(app)
+
+@pytest.fixture
+def client(request):
+    parent_app = Flask(__name__)
+    parent_app.register_blueprint(app, url_prefix=request.param)
+    with parent_app.test_client() as client:
+        yield client
+
+@pytest.mark.parametrize("client, prefix", [(None, ""), ("/prefix","/prefix")], indirect=["client"])
+def test_flask_validate(client, prefix):
+    resp = client.get(prefix+'/ping')
+    assert resp.status_code == 422
+    assert resp.headers.get('X-Error') == 'Validation Error'
+
+    resp = client.get(prefix+'/ping', headers={'lang': 'en-US'})
+    assert resp.json == {'msg': 'pong'}
+    assert resp.headers.get('X-Error') is None
+    assert resp.headers.get('X-Validation') == 'Pass'
+
+    resp = client.post(prefix+'/api/user/flask')
+    assert resp.status_code == 422
+    assert resp.headers.get('X-Error') == 'Validation Error'
+
+    client.set_cookie('flask', 'pub', 'abcdefg')
+    resp = client.post(
+        prefix+'/api/user/flask?order=1',
+        data=json.dumps(dict(name='flask', limit=10)),
+        content_type='application/json',
+    )
+    assert resp.status_code == 200, resp.json
+    assert resp.headers.get('X-Validation') is None
+    assert resp.headers.get('X-API') == 'OK'
+    assert resp.json['name'] == 'flask'
+    assert resp.json['score'] == sorted(resp.json['score'], reverse=True)
+
+    resp = client.post(
+        prefix+'/api/user/flask?order=0',
+        data=json.dumps(dict(name='flask', limit=10)),
+        content_type='application/json',
+    )
+    assert resp.json['score'] == sorted(resp.json['score'], reverse=False)
+
+
+@pytest.mark.parametrize("client, prefix", [(None, ""), ("/prefix","/prefix")], indirect=["client"])
+def test_flask_doc(client, prefix):
+    resp = client.get(prefix+'/apidoc/openapi.json')
+    assert resp.json == api.spec
+
+    resp = client.get(prefix+'/apidoc/redoc')
+    assert resp.status_code == 200
+
+    resp = client.get(prefix+'/apidoc/swagger')
+    assert resp.status_code == 200