Fix yaml rendering

1 files changed, 53 insertions, 0 deletions

diff --git a/rest_framework/utils/encoders.py b/rest_framework/utils/encoders.py
index 6655acbc..2d1fb353 100644
--- a/rest_framework/utils/encoders.py
+++ b/rest_framework/utils/encoders.py
@@ -3,8 +3,11 @@ Helper classes for parsers.
 """
 import datetime
 import decimal
+import types
 from django.utils import simplejson as json
+from django.utils.datastructures import SortedDict
 from rest_framework.compat import timezone
+from rest_framework.serializers import DictWithMetadata, SortedDictWithMetadata
 
 
 class JSONEncoder(json.JSONEncoder):
@@ -36,3 +39,53 @@ class JSONEncoder(json.JSONEncoder):
         elif hasattr(o, '__iter__'):
             return [i for i in o]
         return super(JSONEncoder, self).default(o)
+
+
+try:
+    import yaml
+except ImportError:
+    SafeDumper = None
+else:
+    # Adapted from http://pyyaml.org/attachment/ticket/161/use_ordered_dict.py
+    class SafeDumper(yaml.SafeDumper):
+        """
+        Handles decimals as strings.
+        Handles SortedDicts as usual dicts, but preserves field order, rather
+        than the usual behaviour of sorting the keys.
+        """
+        def represent_decimal(self, data):
+            return self.represent_scalar('tag:yaml.org,2002:str', str(data))
+
+        def represent_mapping(self, tag, mapping, flow_style=None):
+            value = []
+            node = yaml.MappingNode(tag, value, flow_style=flow_style)
+            if self.alias_key is not None:
+                self.represented_objects[self.alias_key] = node
+            best_style = True
+            if hasattr(mapping, 'items'):
+                mapping = list(mapping.items())
+                if not isinstance(mapping, SortedDict):
+                    mapping.sort()
+            for item_key, item_value in mapping:
+                node_key = self.represent_data(item_key)
+                node_value = self.represent_data(item_value)
+                if not (isinstance(node_key, yaml.ScalarNode) and not node_key.style):
+                    best_style = False
+                if not (isinstance(node_value, yaml.ScalarNode) and not node_value.style):
+                    best_style = False
+                value.append((node_key, node_value))
+            if flow_style is None:
+                if self.default_flow_style is not None:
+                    node.flow_style = self.default_flow_style
+                else:
+                    node.flow_style = best_style
+            return node
+
+    SafeDumper.add_representer(SortedDict,
+            yaml.representer.SafeRepresenter.represent_dict)
+    SafeDumper.add_representer(DictWithMetadata,
+            yaml.representer.SafeRepresenter.represent_dict)
+    SafeDumper.add_representer(SortedDictWithMetadata,
+            yaml.representer.SafeRepresenter.represent_dict)
+    SafeDumper.add_representer(types.GeneratorType,
+            yaml.representer.SafeRepresenter.represent_list)