pyyaml custom loader with yaml.safe_load

pyyaml-custom-loader.py

import yaml
from yaml.constructor import ConstructorError

try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader


def no_duplicates_constructor(loader, node, deep=False):
    """Check for duplicate keys."""

    mapping = {}
    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:
            raise ConstructorError("while constructing a mapping", node.start_mark,
                                   "found duplicate key (%s)" % key, key_node.start_mark)
        mapping[key] = value

    return loader.construct_mapping(node, deep)

yaml.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, no_duplicates_constructor, yaml.SafeLoader)

yaml_loaders = [yaml.load, yaml.safe_load]

# Works fine (no duplicate keys)
[yaml_loader('''
---
foo: bar
baz: qux
'''
) for yaml_loader in yaml_loaders]

# Works fine (no duplicate keys on the same level)
[yaml_loader('''
---
foo:
    bar: baz
    baz: qux
bar:
    bar: baz
    baz: qux
'''
) for yaml_loader in yaml_loaders]

# Raises exception (has duplicate keys)
[yaml_loader('''
---
foo: bar
foo: qux
'''
) for yaml_loader in yaml_loaders]