1import yaml
2import yaml.constructor
3
4try:
5    # included in standard lib from Python 2.7
6    from collections import OrderedDict
7except ImportError:
8    # try importing the backported drop-in replacement
9    # it's available on PyPI
10    from ordereddict import OrderedDict
11
12
13class OrderedDictYAMLLoader(yaml.Loader):
14    """
15    A YAML loader that loads mappings into ordered dictionaries
16    (taken from: https://gist.github.com/844388)
17    """
18
19    def __init__(self, *args, **kwargs):
20        yaml.Loader.__init__(self, *args, **kwargs)
21
22        self.add_constructor(u'tag:yaml.org,2002:map', type(self).construct_yaml_map)
23        self.add_constructor(u'tag:yaml.org,2002:omap', type(self).construct_yaml_map)
24
25    def construct_yaml_map(self, node):
26        data = OrderedDict()
27        yield data
28        value = self.construct_mapping(node)
29        data.update(value)
30
31    def construct_mapping(self, node, deep=False):
32        if isinstance(node, yaml.MappingNode):
33            self.flatten_mapping(node)
34        else:
35            raise yaml.constructor.ConstructorError(None, None,
36                'expected a mapping node, but found %s' % node.id, node.start_mark)
37
38        mapping = OrderedDict()
39        for key_node, value_node in node.value:
40            key = self.construct_object(key_node, deep=deep)
41            try:
42                hash(key)
43            except TypeError, exc:
44                raise yaml.constructor.ConstructorError('while constructing a mapping',
45                    node.start_mark, 'found unacceptable key (%s)' % exc, key_node.start_mark)
46            value = self.construct_object(value_node, deep=deep)
47            mapping[key] = value
48        return mapping
49
50if __name__ == '__main__':
51    import textwrap
52
53    sample = """
54one:
55    two: fish
56    red: fish
57    blue: fish
58two:
59    a: yes
60    b: no
61    c: null
62"""
63
64    data = yaml.load(textwrap.dedent(sample), OrderedDictYAMLLoader)
65    assert type(data) is OrderedDict
66    print data