Parse as YAML if not valid JSON

Tests are currently broken, one because invalid JSON
can be valid YAML, and coverage is incomplete.

1 files changed, 119 insertions, 0 deletions

diff --git a/morphlib/yamlparse.py b/morphlib/yamlparse.py
new file mode 100644
index 00000000..7f8b00e5
--- /dev/null
+++ b/morphlib/yamlparse.py
@@ -0,0 +1,119 @@
+# Copyright (C) 2013  Codethink Limited
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+import yaml
+import yaml.constructor
+
+from morphlib.util import OrderedDict
+
+class OrderedDictYAMLLoader(yaml.Loader):
+    """A YAML loader that loads mappings into ordered dictionaries.
+
+       When YAML is loaded with this Loader, it loads mappings as ordered
+       dictionaries, so the order the keys were written in is maintained.
+
+       When combined with the OrderedDictYAMLDumper, this allows yaml documents
+       to be written out in a similar format to they were read.
+
+    """
+
+    def __init__(self, *args, **kwargs):
+        yaml.Loader.__init__(self, *args, **kwargs)
+
+        # When YAML encounters a mapping (which YAML identifies with
+        # the given tag), it will use construct_yaml_map to read it as
+        # an OrderedDict.
+        self.add_constructor(u'tag:yaml.org,2002:map',
+                             type(self).construct_yaml_map)
+
+    def construct_yaml_map(self, node):
+        data = OrderedDict()
+        yield data
+        value = self.construct_mapping(node)
+        data.update(value)
+
+    def construct_mapping(self, node, deep=False):
+        if isinstance(node, yaml.MappingNode):
+            self.flatten_mapping(node)
+        else:
+            raise yaml.constructor.ConstructorError(
+                None, None,
+                'expected a mapping node, but found %s' % node.id,
+                node.start_mark)
+
+        mapping = OrderedDict()
+        for key_node, value_node in node.value:
+            key = self.construct_object(key_node, deep=deep)
+            try:
+                hash(key)
+            except TypeError, exc:
+                raise yaml.constructor.ConstructorError(
+                    'while constructing a mapping', node.start_mark,
+                    'found unacceptable key (%s)' % exc, key_node.start_mark)
+            value = self.construct_object(value_node, deep=deep)
+            mapping[key] = value
+        return mapping
+
+class OrderedDictYAMLDumper(yaml.Dumper):
+    """A YAML dumper that will dump OrderedDicts as mappings.
+
+       When YAML is dumped with this Dumper, it dumps OrderedDicts as
+       mappings, preserving the key order, so the order the keys were
+       written in is maintained.
+
+       When combined with the OrderedDictYAMLDumper, this allows yaml documents
+       to be written out in a similar format to they were read.
+
+    """
+
+    def __init__(self, *args, **kwargs):
+        yaml.Dumper.__init__(self, *args, **kwargs)
+
+        # When YAML sees an OrderedDict, use represent_ordered_dict to dump it
+        self.add_representer(OrderedDict,
+                             type(self).represent_ordered_dict)
+
+    def represent_ordered_dict(self, odict):
+        return self.represent_ordered_mapping(u'tag:yaml.org,2002:map', odict)
+
+    def represent_ordered_mapping(self, tag, omap):
+        value = []
+        node = yaml.MappingNode(tag, value)
+        if self.alias_key is not None:
+            self.represented_objects[self.alias_key] = node
+        best_style = True
+        for item_key, item_value in omap.iteritems():
+            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 # pragma: no cover
+            if not (isinstance(node_value, yaml.ScalarNode) and
+                    not node_value.style):
+                best_style = False # pragma: no cover
+            value.append((node_key, node_value))
+        if self.default_flow_style is not None:
+            node.flow_style = self.default_flow_style
+        else:
+            node.flow_style = best_style # pragma: no cover
+        return node
+
+def load(*args, **kwargs):
+    return yaml.load(Loader=OrderedDictYAMLLoader, *args, **kwargs)
+
+def dump(*args, **kwargs):
+    if 'default_flow_style' not in kwargs:
+        kwargs['default_flow_style'] = False
+    return yaml.dump(Dumper=OrderedDictYAMLDumper, *args, **kwargs)