order_mode support for safe_dump, error handling, and tests

octodns/provider/yaml.py

@@ -430,7 +430,7 @@ class YamlProvider(BaseProvider):
 
                 with open(filename, 'w') as fh:
                     record_data = {record: config}
-                    safe_dump(record_data, fh)
+                    safe_dump(record_data, fh, order_mode=self.order_mode)
 
             if catchall:
                 # Scrub the trailing . to make filenames more sane.
@@ -439,14 +439,19 @@ class YamlProvider(BaseProvider):
                     '_apply:   writing catchall filename=%s', filename
                 )
                 with open(filename, 'w') as fh:
-                    safe_dump(catchall, fh)
+                    safe_dump(catchall, fh, order_mode=self.order_mode)
 
         else:
             # single large file
             filename = join(self.directory, f'{desired.decoded_name}yaml')
             self.log.debug('_apply:   writing filename=%s', filename)
             with open(filename, 'w') as fh:
-                safe_dump(dict(data), fh, allow_unicode=True)
+                safe_dump(
+                    dict(data),
+                    fh,
+                    allow_unicode=True,
+                    order_mode=self.order_mode,
+                )
 
 
 class SplitYamlProvider(YamlProvider):

octodns/yaml.py

@@ -83,12 +83,27 @@ SimpleSortEnforcingLoader.add_constructor(
 )
 
 
+_loaders = {
+    'natural': NaturalSortEnforcingLoader,
+    'simple': SimpleSortEnforcingLoader,
+}
+
+
+class InvalidOrder(Exception):
+
+    def __init__(self, order_mode):
+        options = '", "'.join(_loaders.keys())
+        super().__init__(
+            f'Invalid order_mode, "{order_mode}", options are "{options}"'
+        )
+
+
 def safe_load(stream, enforce_order=True, order_mode='natural'):
     if enforce_order:
-        loader = {
-            'natural': NaturalSortEnforcingLoader,
-            'simple': SimpleSortEnforcingLoader,
-        }[order_mode]
+        try:
+            loader = _loaders[order_mode]
+        except KeyError as e:
+            raise InvalidOrder(order_mode) from e
     else:
         loader = ContextLoader
 
@@ -105,7 +120,7 @@ class SortingDumper(SafeDumper):
     '''
 
     def _representer(self, data):
-        data = sorted(data.items(), key=lambda d: _natsort_key(d[0]))
+        data = sorted(data.items(), key=self.KEYGEN)
         return self.represent_mapping(self.DEFAULT_MAPPING_TAG, data)
 
 
@@ -116,7 +131,18 @@ SortingDumper.add_multi_representer(str, SafeRepresenter.represent_str)
 SortingDumper.add_multi_representer(dict, SortingDumper._representer)
 
 
-def safe_dump(data, fh, **options):
+class NaturalSortingDumper(SortingDumper):
+    KEYGEN = _natsort_key
+
+
+class SimpleSortingDumper(SortingDumper):
+    KEYGEN = lambda _, s: s
+
+
+_dumpers = {'natural': NaturalSortingDumper, 'simple': SimpleSortingDumper}
+
+
+def safe_dump(data, fh, order_mode='natural', **options):
     kwargs = {
         'canonical': False,
         'indent': 2,
@@ -125,4 +151,8 @@ def safe_dump(data, fh, **options):
         'explicit_start': True,
     }
     kwargs.update(options)
-    dump(data, fh, SortingDumper, **kwargs)
+    try:
+        dumper = _dumpers[order_mode]
+    except KeyError as e:
+        raise InvalidOrder(order_mode) from e
+    dump(data, fh, dumper, **kwargs)

tests/test_octodns_yaml.py

@@ -7,7 +7,7 @@ from unittest import TestCase
 
 from yaml.constructor import ConstructorError
 
-from octodns.yaml import safe_dump, safe_load
+from octodns.yaml import InvalidOrder, safe_dump, safe_load
 
 
 class TestYaml(TestCase):
@@ -88,30 +88,58 @@ class TestYaml(TestCase):
             )
 
     def test_order_mode(self):
+        data = {'*.1.2': 'a', '*.10.1': 'c', '*.11.2': 'd', '*.2.2': 'b'}
+        natural = '''---
+'*.1.2': a
+'*.2.2': b
+'*.10.1': c
+'*.11.2': d
+'''
+        simple = '''---
+'*.1.2': a
+'*.10.1': c
+'*.11.2': d
+'*.2.2': b
+'''
+
+        ## natural
+        # correct order
+        self.assertEqual(data, safe_load(natural))
+        # wrong order
+        with self.assertRaises(ConstructorError) as ctx:
+            safe_load(simple)
+        problem = ctx.exception.problem.split(' at')[0]
         self.assertEqual(
-            {'*.1.2': 'a', '*.10.1': 'c', '*.11.2': 'd', '*.2.2': 'b'},
-            safe_load(
-                '''
-'*.1.2': 'a'
-'*.10.1': 'c'
-'*.11.2': 'd'
-'*.2.2': 'b'
-''',
-                order_mode='simple',
-            ),
+            'keys out of order: expected *.2.2 got *.10.1', problem
         )
-        # natural sort throws error
+        # dump
+        buf = StringIO()
+        safe_dump(data, buf)
+        self.assertEqual(natural, buf.getvalue())
+
+        ## simple
+        # correct order
+        self.assertEqual(data, safe_load(simple, order_mode='simple'))
+        # wrong order
         with self.assertRaises(ConstructorError) as ctx:
-            safe_load(
-                '''
-'*.1.2': 'a'
-'*.2.2': 'b'
-'*.10.1': 'c'
-'*.11.2': 'd'
-''',
-                order_mode='simple',
-            )
-        self.assertTrue(
-            'keys out of order: expected *.10.1 got *.2.2 at'
-            in ctx.exception.problem
+            safe_load(natural, order_mode='simple')
+        problem = ctx.exception.problem.split(' at')[0]
+        self.assertEqual(
+            'keys out of order: expected *.10.1 got *.2.2', problem
+        )
+        buf = StringIO()
+        safe_dump(data, buf, order_mode='simple')
+        self.assertEqual(simple, buf.getvalue())
+
+        with self.assertRaises(InvalidOrder) as ctx:
+            safe_load(None, order_mode='bad')
+        self.assertEqual(
+            'Invalid order_mode, "bad", options are "natural", "simple"',
+            str(ctx.exception),
+        )
+        with self.assertRaises(InvalidOrder) as ctx:
+            safe_dump(None, None, order_mode='bad2')
+        self.assertEqual(
+            'Invalid order_mode, "bad2", options are "natural", "simple"',
+            str(ctx.exception),
         )