From d5f399e09c0c836a7abd61e80abb66b0247457cd Mon Sep 17 00:00:00 2001 From: Ross McFarland Date: Tue, 8 Oct 2024 10:28:33 -0700 Subject: [PATCH] order_mode support for safe_dump, error handling, and tests --- octodns/provider/yaml.py | 11 ++++-- octodns/yaml.py | 44 ++++++++++++++++++---- tests/test_octodns_yaml.py | 76 ++++++++++++++++++++++++++------------ 3 files changed, 97 insertions(+), 34 deletions(-) diff --git a/octodns/provider/yaml.py b/octodns/provider/yaml.py index e4afffe..b5fd3d7 100644 --- a/octodns/provider/yaml.py +++ b/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): diff --git a/octodns/yaml.py b/octodns/yaml.py index 0368e63..5638df5 100644 --- a/octodns/yaml.py +++ b/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) diff --git a/tests/test_octodns_yaml.py b/tests/test_octodns_yaml.py index 01718c7..a05d28d 100644 --- a/tests/test_octodns_yaml.py +++ b/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), )