Ticket #161: use_ordered_dict.py

File use_ordered_dict.py, 2.2 KB (added by anonymous, 4 years ago)
Line 
1import yaml
2import collections
3
4def construct_ordered_mapping(self, node, deep=False):
5    if not isinstance(node, yaml.MappingNode):
6        raise ConstructorError(None, None,
7                "expected a mapping node, but found %s" % node.id,
8                node.start_mark)
9    mapping = collections.OrderedDict()
10    for key_node, value_node in node.value:
11        key = self.construct_object(key_node, deep=deep)
12        if not isinstance(key, collections.Hashable):
13            raise ConstructorError("while constructing a mapping", node.start_mark,
14                    "found unhashable key", key_node.start_mark)
15        value = self.construct_object(value_node, deep=deep)
16        mapping[key] = value
17    return mapping
18
19yaml.constructor.BaseConstructor.construct_mapping = construct_ordered_mapping
20
21
22def construct_yaml_map_with_ordered_dict(self, node):
23    data = collections.OrderedDict()
24    yield data
25    value = self.construct_mapping(node)
26    data.update(value)
27
28yaml.constructor.Constructor.add_constructor(
29        'tag:yaml.org,2002:map',
30        construct_yaml_map_with_ordered_dict)
31
32
33def represent_ordered_mapping(self, tag, mapping, flow_style=None):
34    value = []
35    node = yaml.MappingNode(tag, value, flow_style=flow_style)
36    if self.alias_key is not None:
37        self.represented_objects[self.alias_key] = node
38    best_style = True
39    if hasattr(mapping, 'items'):
40        mapping = list(mapping.items())
41    for item_key, item_value in mapping:
42        node_key = self.represent_data(item_key)
43        node_value = self.represent_data(item_value)
44        if not (isinstance(node_key, yaml.ScalarNode) and not node_key.style):
45            best_style = False
46        if not (isinstance(node_value, yaml.ScalarNode) and not node_value.style):
47            best_style = False
48        value.append((node_key, node_value))
49    if flow_style is None:
50        if self.default_flow_style is not None:
51            node.flow_style = self.default_flow_style
52        else:
53            node.flow_style = best_style
54    return node
55
56yaml.representer.BaseRepresenter.represent_mapping = represent_ordered_mapping
57
58yaml.representer.Representer.add_representer(collections.OrderedDict,
59        yaml.representer.SafeRepresenter.represent_dict)