Merge pull request flask-restful#214 from lyschoening/marshal-with-field

Added marshal_with_field decorator, see issue flask-restful#212

Author: dougblack

docs/api.rst

@@ -8,6 +8,7 @@ API Docs
 
 .. autofunction:: marshal
 .. autofunction:: marshal_with
+.. autofunction:: marshal_with_field
 .. autofunction:: abort
 
 

flask_restful/__init__.py

@@ -21,7 +21,7 @@
 except ImportError:
     from .utils.ordereddict import OrderedDict
 
-__all__ = ('Api', 'Resource', 'marshal', 'marshal_with', 'abort')
+__all__ = ('Api', 'Resource', 'marshal', 'marshal_with', 'marshal_with_field', 'abort')
 
 
 def abort(http_status_code, **kwargs):
@@ -562,3 +562,39 @@ def wrapper(*args, **kwargs):
             else:
                 return marshal(resp, self.fields)
         return wrapper
+
+
+class marshal_with_field(object):
+    """
+    A decorator that formats the return values of your methods with a single field.
+
+    >>> from flask.ext.restful import marshal_with_field, fields
+    >>> @marshal_with_field(fields.List(fields.Integer))
+    ... def get():
+    ...     return ['1', 2, 3.0]
+    ...
+    >>> get()
+    [1, 2, 3]
+
+    see :meth:`flask.ext.restful.marshal_with`
+    """
+    def __init__(self, field):
+        """
+        :param field: a single field with which to marshal the output.
+        """
+        if isinstance(field, type):
+            self.field = field()
+        else:
+            self.field = field
+
+    def __call__(self, f):
+        @wraps(f)
+        def wrapper(*args, **kwargs):
+            resp = f(*args, **kwargs)
+
+            if isinstance(resp, tuple):
+                data, code, headers = unpack(resp)
+                return self.field.format(data), code, headers
+            return self.field.format(resp)
+
+        return wrapper

flask_restful/fields.py

@@ -145,22 +145,25 @@ def __init__(self, cls_or_instance, **kwargs):
                 raise MarshallingException(error_msg)
             self.container = cls_or_instance
 
+    def format(self, value):
+        # Convert all instances in typed list to container type
+        if isinstance(value, set):
+            value = list(value)
+
+        return [
+            self.container.output(idx,
+                                  val if isinstance(val, dict) and
+                                         not isinstance(self.container, Nested)
+                                  and not type(self.container) is Raw
+                                  else value)
+            for idx, val in enumerate(value)
+        ]
+
     def output(self, key, data):
         value = get_value(key if self.attribute is None else self.attribute, data)
         # we cannot really test for external dict behavior
         if is_indexable_but_not_string(value) and not isinstance(value, dict):
-            # Convert all instances in typed list to container type
-            if isinstance(value, set):
-                value = list(value)
-
-            return [
-                self.container.output(idx,
-                                      val if isinstance(val, dict) and
-                                      not isinstance(self.container, Nested)
-                                      and not type(self.container) is Raw
-                                      else value)
-                for idx, val in enumerate(value)
-            ]
+            return self.format(value)
 
         if value is None:
             return self.default

tests/test_api.py

@@ -117,6 +117,22 @@ def try_me():
             return OrderedDict([('foo', 'bar'), ('bat', 'baz')]), 200, {'X-test': 123}
         self.assertEquals(try_me(), ({'foo': 'bar'}, 200, {'X-test': 123}))
 
+    def test_marshal_field_decorator(self):
+        field = flask_restful.fields.Raw
+
+        @flask_restful.marshal_with_field(field)
+        def try_me():
+            return 'foo'
+        self.assertEquals(try_me(), 'foo')
+
+    def test_marshal_field_decorator_tuple(self):
+        field = flask_restful.fields.Raw
+
+        @flask_restful.marshal_with_field(field)
+        def try_me():
+            return 'foo', 200, {'X-test': 123}
+        self.assertEquals(('foo', 200, {'X-test': 123}), try_me())
+
     def test_marshal_field(self):
         fields = OrderedDict({'foo': flask_restful.fields.Raw()})
         marshal_fields = OrderedDict([('foo', 'bar'), ('bat', 'baz')])