Add callable sort key functionality to sort_keys

lib/yaml/representer.py

@@ -110,7 +110,10 @@ def represent_mapping(self, tag, mapping, flow_style=None):
             mapping = list(mapping.items())
             if self.sort_keys:
                 try:
-                    mapping = sorted(mapping)
+                    if callable(self.sort_keys):
+                        mapping = sorted(mapping, key=self.sort_keys)
+                    else:
+                        mapping = sorted(mapping)
                 except TypeError:
                     pass
         for item_key, item_value in mapping:

tests/data/mapping.sort_vowels_first

@@ -0,0 +1,6 @@
+u: 1
+c: 2
+v: 3
+d: 4
+w: 5
+e: 6

tests/data/mapping.sorted_vowels_first

@@ -0,0 +1,6 @@
+e: 6
+u: 1
+c: 2
+d: 4
+v: 3
+w: 5

tests/lib/test_sort_keys_function.py

@@ -0,0 +1,41 @@
+import yaml
+import sys
+import functools
+
+_vowels = set(('a','e','i','o','u'))
+
+def sort_keys_vowels_first(a, b):
+    # params are (key,value) pairs
+    a=a[0]
+    b=b[0]
+    if (a in _vowels) == (b in _vowels):
+        return (a>b)-(a<b) # cmp(a,b)
+    else:
+        return (b in _vowels)*2-1
+
+def test_sort_keys_function(input_filename, sorted_filename, verbose=False):
+    with open(input_filename, 'rb') as file:
+        input = file.read().decode('utf-8')
+    with open(sorted_filename, 'rb') as file:
+        sorted = file.read().decode('utf-8')
+    data = yaml.load(input, Loader=yaml.FullLoader)
+    dump_sorted = yaml.dump(data, default_flow_style=False, sort_keys=functools.cmp_to_key(sort_keys_vowels_first))
+    dump_unsorted = yaml.dump(data, default_flow_style=False, sort_keys=False)
+    dump_unsorted_safe = yaml.dump(data, default_flow_style=False, sort_keys=False, Dumper=yaml.SafeDumper)
+    if verbose:
+        print("INPUT:")
+        print(input)
+        print("DATA:")
+        print(data)
+
+    assert dump_sorted == sorted
+
+    if sys.version_info>=(3,7):
+        assert dump_unsorted == input
+        assert dump_unsorted_safe == input
+
+test_sort_keys_function.unittest = ['.sort_vowels_first', '.sorted_vowels_first']
+
+if __name__ == '__main__':
+    import test_appliance
+    test_appliance.run(globals())

tests/lib/test_yaml.py

@@ -13,6 +13,7 @@
 from test_recursive import *
 from test_input_output import *
 from test_sort_keys import *
+from test_sort_keys_function import *
 from test_multi_constructor import *
 
 from test_schema import *