Getting duplicate keys in YAML using Python

Viewed 14316

We are in need of parsing YAML files which contain duplicate keys and all of these need to be parsed. It is not enough to skip duplicates. I know this is against the YAML spec and I would like to not have to do it, but a third-party tool used by us enables this usage and we need to deal with it.

File example:

build:
  step: 'step1'

build:
  step: 'step2'

After parsing we should have a similar data structure to this:

yaml.load('file.yml')
# [('build', [('step', 'step1')]), ('build', [('step', 'step2')])]

dict can no longer be used to represent the parsed contents.

I am looking for a solution in Python and I didn't find a library supporting this, have I missed anything?

Alternatively, I am happy to write my own thing but would like to make it as simple as possible. ruamel.yaml looks like the most advanced YAML parser in Python and it looks moderately extensible, can it be extended to support duplicate fields?

4 Answers

You can override how pyyaml loads keys. For example, you could use a defaultdict with lists of values for each keys:

from collections import defaultdict
import yaml


def parse_preserving_duplicates(src):
    # We deliberately define a fresh class inside the function,
    # because add_constructor is a class method and we don't want to
    # mutate pyyaml classes.
    class PreserveDuplicatesLoader(yaml.loader.Loader):
        pass

    def map_constructor(loader, node, deep=False):
        """Walk the mapping, recording any duplicate keys.

        """
        mapping = defaultdict(list)
        for key_node, value_node in node.value:
            key = loader.construct_object(key_node, deep=deep)
            value = loader.construct_object(value_node, deep=deep)

            mapping[key].append(value)

        return mapping

    PreserveDuplicatesLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, map_constructor)
    return yaml.load(src, PreserveDuplicatesLoader)

Here is an alternative implementation based on Anthon's answer and ruamel.yaml. It is rather generic and uses lists for duplicates, while other entries are left unchanged.

from collections import Counter
from ruamel.yaml import YAML
from ruamel.yaml.constructor import SafeConstructor

yaml_str = '''
a: 1
b: 2
b: 2
'''

def construct_yaml_map(self, node):
    data = {}
    yield data
    keys = [self.construct_object(node, deep=True) for node, _ in node.value]
    vals = [self.construct_object(node, deep=True) for _, node in node.value]
    key_count = Counter(keys)
    for key, val in zip(keys, vals):
        if key_count[key] > 1:
            if key not in data:
                data[key] = []
            data[key].append(val)
        else:
            data[key] = val

SafeConstructor.add_constructor(u'tag:yaml.org,2002:map', construct_yaml_map)
yaml = YAML(typ='safe')
data = yaml.load(yaml_str)
print(data)

Output:

{'a': 1, 'b': [2, 2]}

The same is possible with the pyyaml package (inspired by Wilfred Hughes' answer):

from collections import Counter
import yaml

yaml_str = '''
a: 1
b: 2
b: 2
'''

def parse_preserving_duplicates(src):
    class PreserveDuplicatesLoader(yaml.loader.Loader):
        pass

    def map_constructor(loader, node, deep=False):
        keys = [loader.construct_object(node, deep=deep) for node, _ in node.value]
        vals = [loader.construct_object(node, deep=deep) for _, node in node.value]
        key_count = Counter(keys)
        data = {}
        for key, val in zip(keys, vals):
            if key_count[key] > 1:
                if key not in data:
                    data[key] = []
                data[key].append(val)
            else:
                data[key] = val
        return data

    PreserveDuplicatesLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, map_constructor)
    return yaml.load(src, PreserveDuplicatesLoader)

print(parse_preserving_duplicates(yaml_str))

Output:

{'a': 1, 'b': [2, 2]}
Related