diff --git a/pre_commit/yaml_extensions.py b/pre_commit/yaml_extensions.py new file mode 100644 index 00000000..70bcec72 --- /dev/null +++ b/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) diff --git a/tests/yaml_extensions_test.py b/tests/yaml_extensions_test.py new file mode 100644 index 00000000..c60c1aef --- /dev/null +++ b/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' + )