diff --git a/bin/compose-core b/bin/compose-core index f081d31..cf61c93 100755 --- a/bin/compose-core +++ b/bin/compose-core @@ -134,6 +134,16 @@ export _merge_yaml_common_code=" import sys import yaml +try: + from yaml import CSafeLoader as SafeLoader, CSafeDumper as SafeDumper +except ImportError: ## pragma: no cover + sys.stderr.write('YAML code in pure python\n') + exit(1) + from yaml import SafeLoader, SafeDumper + +class MySafeLoader(SafeLoader): pass +class MySafeDumper(SafeDumper): pass + try: # included in standard lib from Python 2.7 @@ -149,13 +159,50 @@ except ImportError: class MyOrderedDict(OrderedDict): pass -yaml.add_representer( +MySafeDumper.add_representer( MyOrderedDict, lambda cls, data: cls.represent_dict(data.items())) -yaml.add_constructor( + +def construct_omap(cls, node): + cls.flatten_mapping(node) + return MyOrderedDict(cls.construct_pairs(node)) + + +MySafeLoader.add_constructor( yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, - lambda cls, node: MyOrderedDict(cls.construct_pairs(node))) + construct_omap) + + +## +## Support local and global objects +## + +class EncapsulatedNode(object): pass + + +def mk_encapsulated_node(s, node): + + method = 'construct_%s' % (node.id, ) + data = getattr(s, method)(node) + + class _E(data.__class__, EncapsulatedNode): + pass + + _E.__name__ = str(node.tag) + _E._node = node + return _E(data) + + +def represent_encapsulated_node(s, o): + value = s.represent_data(o.__class__.__bases__[0](o)) + value.tag = o.__class__.__name__ + return value + + +MySafeDumper.add_multi_representer(EncapsulatedNode, + represent_encapsulated_node) +MySafeLoader.add_constructor(None, mk_encapsulated_node) def fc(filename): @@ -218,7 +265,7 @@ def merge_cli(*args): for i, a in enumerate(e.args[2])))) exit(1) if c is not None: - print '%s' % yaml.dump(c, default_flow_style=False) + print '%s' % yaml.dump(c, default_flow_style=False, Dumper=MySafeDumper) " @@ -230,7 +277,7 @@ merge_yaml() { $_merge_yaml_common_code -merge_cli(*(yaml.load(fc(f)) for f in sys.argv[1:])) +merge_cli(*(yaml.load(fc(f), Loader=MySafeLoader) for f in sys.argv[1:])) EOF fi @@ -247,7 +294,7 @@ merge_yaml_str() { $_merge_yaml_common_code -merge_cli(*(yaml.load(f) for f in sys.argv[1:])) +merge_cli(*(yaml.load(f, Loader=MySafeLoader) for f in sys.argv[1:])) EOF fi @@ -264,9 +311,14 @@ yaml_key_val_str() { $_merge_yaml_common_code -print '%s' % yaml.dump({ - yaml.load(sys.argv[1]): - yaml.load(sys.argv[2])}, default_flow_style=False) +print '%s' % yaml.dump( + { + yaml.load(sys.argv[1], Loader=MySafeLoader): + yaml.load(sys.argv[2], Loader=MySafeLoader) + }, + default_flow_style=False, + Dumper=MySafeDumper, + ) EOF fi