Source code for galaxy.util.yaml_util

from __future__ import absolute_import

import logging
import os
from collections import OrderedDict

import yaml
try:
    from yaml import CSafeLoader as SafeLoader
except ImportError:
    from yaml import SafeLoader
from yaml.constructor import ConstructorError


log = logging.getLogger(__name__)


[docs]class OrderedLoader(SafeLoader):
    # This class was pulled out of ordered_load() for the sake of
    # mocking __init__ in a unit test.
[docs]    def __init__(self, stream):
        self._root = os.path.split(stream.name)[0]
        super(OrderedLoader, self).__init__(stream)

[docs]    def include(self, node):
        filename = os.path.join(self._root, self.construct_scalar(node))
        with open(filename, 'r') as f:
            return yaml.load(f, OrderedLoader)


[docs]def ordered_load(stream, merge_duplicate_keys=False):
    """
    Parse the first YAML document in a stream and produce the corresponding
    Python object, using OrderedDicts instead of dicts.

    If merge_duplicate_keys is True, merge the values of duplicate mapping keys
    into a list, as the uWSGI "dumb" YAML parser would do.
    Otherwise, following YAML 1.2 specification which says that "each key is
    unique in the association", raise a ConstructionError exception.
    """
    def construct_mapping(loader, node, deep=False):
        loader.flatten_mapping(node)
        mapping = OrderedDict()
        merged_duplicate = {}
        for key_node, value_node in node.value:
            key = loader.construct_object(key_node, deep=deep)
            value = loader.construct_object(value_node, deep=deep)
            if key in mapping:
                if not merge_duplicate_keys:
                    raise ConstructorError("while constructing a mapping", node.start_mark,
                        "found duplicated key (%s)" % key, key_node.start_mark)
                log.debug("Merging values for duplicate key '%s' into a list", key)
                if merged_duplicate.get(key):
                    mapping[key].append(value)
                else:
                    mapping[key] = [mapping[key], value]
                    merged_duplicate[key] = True
            else:
                mapping[key] = value
        return mapping

    OrderedLoader.add_constructor(
        yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
        construct_mapping)
    OrderedLoader.add_constructor('!include', OrderedLoader.include)

    return yaml.load(stream, OrderedLoader)


[docs]def ordered_dump(data, stream=None, Dumper=yaml.Dumper, **kwds):
    class OrderedDumper(Dumper):
        pass

    def _dict_representer(dumper, data):
        return dumper.represent_mapping(
            yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
            list(data.items()))
    OrderedDumper.add_representer(OrderedDict, _dict_representer)
    return yaml.dump(data, stream, OrderedDumper, **kwds)