diff --git a/tts_yaml_unpacker.py b/tts_yaml_unpacker.py index 03af549..2627fe9 100755 --- a/tts_yaml_unpacker.py +++ b/tts_yaml_unpacker.py @@ -5,6 +5,14 @@ import yaml import os import sys from pathvalidate import sanitize_filename +from yaml import CDumper as Dumper + + +class Loader(yaml.CLoader): + """Hack to store the stream before passing it off to :class:`yaml.CLoader`""" + def __init__(self, stream): + self.stream = stream + super().__init__(stream) class IncludeTag: @@ -23,13 +31,13 @@ class IncludeTag: target = os.path.join(base_dir, node.value) with open(target) as f: if target.endswith('.yaml'): - return yaml.load(f, Loader=yaml.Loader) + return yaml.load(f, Loader=Loader) else: return f.read() -yaml.add_constructor('!include', IncludeTag.from_yaml) -yaml.add_representer(IncludeTag, IncludeTag.to_yaml) +Loader.add_constructor('!include', IncludeTag.from_yaml) +Dumper.add_representer(IncludeTag, IncludeTag.to_yaml) def uniqueName(obj): @@ -88,7 +96,7 @@ def recursivelyUnpackObject(parent_dir, obj, obj[k] = recursively_normalize_values(obj[k]) with open(file_base_path + '.yaml', 'w') as f: - yaml.dump(obj, f) + yaml.dump(obj, f, Dumper=Dumper) return obj_base_name + '.yaml' @@ -107,7 +115,7 @@ def unpackJson(json_file, output_name): def packYaml(yaml_file, output_json_file): with open(yaml_file) as f: - data = yaml.load(f, Loader=yaml.Loader) + data = yaml.load(f, Loader=Loader) with open(output_json_file, 'w') as f: json.dump(data, f, indent=2)