1
0
mirror of https://github.com/github/octodns.git synced 2024-05-11 05:55:00 +00:00
Files
github-octodns/octodns/yaml.py
2023-01-22 13:54:19 -08:00

73 lines
2.1 KiB
Python

#
#
#
from natsort import natsort_keygen
from yaml import SafeDumper, SafeLoader, dump, load
from yaml.constructor import ConstructorError
from yaml.representer import SafeRepresenter
_natsort_key = natsort_keygen()
# Found http://stackoverflow.com/a/21912744 which guided me on how to hook in
# here
class SortEnforcingLoader(SafeLoader):
def _construct(self, node):
self.flatten_mapping(node)
ret = self.construct_pairs(node)
keys = [d[0] for d in ret]
keys_sorted = sorted(keys, key=_natsort_key)
for key in keys:
expected = keys_sorted.pop(0)
if key != expected:
raise ConstructorError(
None,
None,
'keys out of order: '
f'expected {expected} got {key} at ' + str(node.start_mark),
)
return dict(ret)
SortEnforcingLoader.add_constructor(
SortEnforcingLoader.DEFAULT_MAPPING_TAG, SortEnforcingLoader._construct
)
def safe_load(stream, enforce_order=True):
return load(stream, SortEnforcingLoader if enforce_order else SafeLoader)
class SortingDumper(SafeDumper):
'''
This sorts keys alphanumerically in a "natural" manner where things with
the number 2 come before the number 12.
See https://www.xormedia.com/natural-sort-order-with-zero-padding/ for
more info
'''
def _representer(self, data):
data = sorted(data.items(), key=lambda d: _natsort_key(d[0]))
return self.represent_mapping(self.DEFAULT_MAPPING_TAG, data)
SortingDumper.add_representer(dict, SortingDumper._representer)
# This should handle all the record value types which are ultimately either str
# or dict at some point in their inheritance hierarchy
SortingDumper.add_multi_representer(str, SafeRepresenter.represent_str)
SortingDumper.add_multi_representer(dict, SortingDumper._representer)
def safe_dump(data, fh, **options):
kwargs = {
'canonical': False,
'indent': 2,
'default_style': '',
'default_flow_style': False,
'explicit_start': True,
}
kwargs.update(options)
dump(data, fh, SortingDumper, **kwargs)