pytorch

1
# Safely load fast C Yaml loader/dumper if they are available
2
try:
3
    from yaml import CSafeLoader as Loader
4
except ImportError:
5
    from yaml import SafeLoader as Loader  # type: ignore[assignment, misc]
6

7
try:
8
    from yaml import CSafeDumper as Dumper
9
except ImportError:
10
    from yaml import SafeDumper as Dumper  # type: ignore[assignment, misc]
11
YamlDumper = Dumper
12

13

14
# A custom loader for YAML that errors on duplicate keys.
15
# This doesn't happen by default: see https://github.com/yaml/pyyaml/issues/165
16
class YamlLoader(Loader):
17
    def construct_mapping(self, node, deep=False):  # type: ignore[no-untyped-def]
18
        mapping = []
19
        for key_node, value_node in node.value:
20
            key = self.construct_object(key_node, deep=deep)  # type: ignore[no-untyped-call]
21
            assert (
22
                key not in mapping
23
            ), f"Found a duplicate key in the yaml. key={key}, line={node.start_mark.line}"
24
            mapping.append(key)
25
        mapping = super().construct_mapping(node, deep=deep)  # type: ignore[no-untyped-call]
26
        return mapping
27