Closes #2

jpadilla · Dec 18, 2014 · df76291 · df76291
1 parent 6da39ee
commit df76291
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/rest_framework_yaml/encoders.py b/rest_framework_yaml/encoders.py
@@ -7,15 +7,19 @@
 
 from django.utils import six
 
+# Note: ReturnDict and ReturnList are private API
+from rest_framework.utils.serializer_helpers import ReturnDict, ReturnList
+
 from .compat import yaml, OrderedDict
 
 
-# Adapted from http://pyyaml.org/attachment/ticket/161/use_ordered_dict.py
 class SafeDumper(yaml.SafeDumper):
     """
     Handles decimals as strings.
     Handles OrderedDicts as usual dicts, but preserves field order, rather
     than the usual behaviour of sorting the keys.
+
+    Adapted from http://pyyaml.org/attachment/ticket/161/use_ordered_dict.py
     """
     def represent_decimal(self, data):
         return self.represent_scalar('tag:yaml.org,2002:str', six.text_type(data))
@@ -55,6 +59,16 @@ def represent_mapping(self, tag, mapping, flow_style=None):
     yaml.representer.SafeRepresenter.represent_dict
 )
 
+SafeDumper.add_representer(
+    ReturnDict,
+    yaml.representer.SafeRepresenter.represent_dict
+)
+
+SafeDumper.add_representer(
+    ReturnList,
+    yaml.representer.SafeRepresenter.represent_list
+)
+
 SafeDumper.add_representer(
     types.GeneratorType,
     yaml.representer.SafeRepresenter.represent_list