Add ordered yaml load/dump

pre_commit/yaml_extensions.py

@@ -0,0 +1,28 @@
+
+import yaml
+
+from pre_commit.ordereddict import OrderedDict
+
+
+# Adapted from http://stackoverflow.com/a/21912744/812183
+
+def ordered_load(s):
+    class OrderedLoader(yaml.loader.Loader): pass
+    def constructor(loader, node):
+        return OrderedDict(loader.construct_pairs(node))
+    OrderedLoader.add_constructor(
+        yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
+        constructor,
+    )
+    return yaml.load(s, Loader=OrderedLoader)
+
+
+def ordered_dump(s, **kwargs):
+    class OrderedDumper(yaml.dumper.Dumper): pass
+    def dict_representer(dumper, data):
+        return dumper.represent_mapping(
+            yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
+            data.items(),
+        )
+    OrderedDumper.add_representer(OrderedDict, dict_representer)
+    return yaml.dump(s, Dumper=OrderedDumper, **kwargs)

tests/yaml_extensions_test.py

@@ -0,0 +1,35 @@
+
+from pre_commit.ordereddict import OrderedDict
+from pre_commit.yaml_extensions import ordered_dump
+from pre_commit.yaml_extensions import ordered_load
+
+
+def test_ordered_load():
+    ret = ordered_load(
+        'a: herp\n'
+        'c: derp\n'
+        'd: darp\n'
+        'b: harp\n'
+    )
+    # Original behavior
+    assert ret == {'a': 'herp', 'b': 'harp', 'c': 'derp', 'd': 'darp'}
+    # Ordered behavior
+    assert (
+        ret.items() ==
+        [('a', 'herp'), ('c', 'derp'), ('d', 'darp'), ('b', 'harp')]
+    )
+
+
+def test_ordered_dump():
+    ret = ordered_dump(
+        OrderedDict(
+            (('a', 'herp'), ('c', 'derp'), ('b', 'harp'), ('d', 'darp'))
+        ),
+        default_flow_style=False,
+    )
+    assert ret == (
+        'a: herp\n'
+        'c: derp\n'
+        'b: harp\n'
+        'd: darp\n'
+    )