Support decorators using either postional args or a single list/tuple

flask-api · Feb 27, 2014 · 0c6576b · 0c6576b
1 parent e16fdcb
commit 0c6576b
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 8 deletions.
diff --git a/flask_api/app.py b/flask_api/app.py
@@ -18,6 +18,7 @@
     template_folder='templates', static_folder='static'
 )
 
+
 def urlize_quoted_links(content):
     return re.sub(r'"(https?://[^"]*)"', r'"<a href="\1">\1</a>"', content)
 

diff --git a/flask_api/decorators.py b/flask_api/decorators.py
@@ -2,21 +2,27 @@
 from flask import request
 
 
-def set_parsers(parsers=None):
+def set_parsers(*parsers):
     def decorator(func):
         @wraps(func)
         def decorated_function(*args, **kwargs):
-            request.parser_classes = parsers
+            if len(parsers) == 1 and isinstance(parsers, (list, tuple)):
+                request.parser_classes = parsers[0]
+            else:
+                request.parser_classes = parsers
             return func(*args, **kwargs)
         return decorated_function
     return decorator
 
 
-def set_renderers(renderers=None):
+def set_renderers(*renderers):
     def decorator(func):
         @wraps(func)
         def decorated_function(*args, **kwargs):
-            request.renderer_classes = renderers
+            if len(renderers) == 1 and isinstance(renderers, (list, tuple)):
+                request.renderer_classes = renderers[0]
+            else:
+                request.renderer_classes = renderers
             return func(*args, **kwargs)
         return decorated_function
     return decorator
diff --git a/flask_api/tests/test_parsers.py b/flask_api/tests/test_parsers.py
@@ -136,6 +136,11 @@ def custom_parser_1():
         def custom_parser_2():
             return {'data': request.data}
 
+        @app.route('/custom_parser_2_as_args/', methods=['POST'])
+        @set_parsers(CustomParser2, CustomParser1)
+        def custom_parser_2_as_args():
+            return {'data': request.data}
+
         self.app = app
 
     def test_overridden_parsers_with_settings(self):
@@ -161,3 +166,15 @@ def test_overridden_parsers_with_decorator(self):
                 "data": "custom parser 2",
             }
             self.assertEqual(data, expected)
+
+    def test_overridden_parsers_with_decorator_as_args(self):
+        with self.app.test_client() as client:
+            data = {'example': 'example'}
+            response = client.post('/custom_parser_2_as_args/', data=data)
+            self.assertEqual(response.status_code, status.HTTP_200_OK)
+            self.assertEqual(response.headers['Content-Type'], 'application/json')
+            data = json.loads(response.get_data().decode('utf8'))
+            expected = {
+                "data": "custom parser 2",
+            }
+            self.assertEqual(data, expected)
diff --git a/flask_api/tests/test_renderers.py b/flask_api/tests/test_renderers.py
@@ -46,20 +46,25 @@ def render(self, data, media_type, **options):
         app.config['DEFAULT_RENDERERS'] = [CustomRenderer1]
         app.config['PROPAGATE_EXCEPTIONS'] = True
 
-        @app.route('/custom_parser_1/', methods=['GET'])
+        @app.route('/custom_renderer_1/', methods=['GET'])
         def custom_renderer_1():
             return {'data': 'example'}
 
-        @app.route('/custom_parser_2/', methods=['GET'])
+        @app.route('/custom_renderer_2/', methods=['GET'])
         @set_renderers([CustomRenderer2])
         def custom_renderer_2():
             return {'data': 'example'}
 
+        @app.route('/custom_renderer_2_as_args/', methods=['GET'])
+        @set_renderers(CustomRenderer2, CustomRenderer1)
+        def custom_renderer_2_as_args():
+            return {'data': 'example'}
+
         self.app = app
 
     def test_overridden_parsers_with_settings(self):
         with self.app.test_client() as client:
-            response = client.get('/custom_parser_1/')
+            response = client.get('/custom_renderer_1/')
             self.assertEqual(response.status_code, status.HTTP_200_OK)
             self.assertEqual(response.headers['Content-Type'], 'application/example1')
             data = response.get_data().decode('utf8')
@@ -68,7 +73,16 @@ def test_overridden_parsers_with_settings(self):
     def test_overridden_parsers_with_decorator(self):
         with self.app.test_client() as client:
             data = {'example': 'example'}
-            response = client.get('/custom_parser_2/', data=data)
+            response = client.get('/custom_renderer_2/', data=data)
+            self.assertEqual(response.status_code, status.HTTP_200_OK)
+            self.assertEqual(response.headers['Content-Type'], 'application/example2')
+            data = response.get_data().decode('utf8')
+            self.assertEqual(data, "custom renderer 2")
+
+    def test_overridden_parsers_with_decorator_as_args(self):
+        with self.app.test_client() as client:
+            data = {'example': 'example'}
+            response = client.get('/custom_renderer_2_as_args/', data=data)
             self.assertEqual(response.status_code, status.HTTP_200_OK)
             self.assertEqual(response.headers['Content-Type'], 'application/example2')
             data = response.get_data().decode('utf8')