xref: /aosp_15_r20/external/pytorch/torchgen/yaml_utils.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# Safely load fast C Yaml loader/dumper if they are available
2try:
3    from yaml import CSafeLoader as Loader
4except ImportError:
5    from yaml import SafeLoader as Loader  # type: ignore[assignment, misc]
6
7try:
8    from yaml import CSafeDumper as Dumper
9except ImportError:
10    from yaml import SafeDumper as Dumper  # type: ignore[assignment, misc]
11YamlDumper = 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
16class 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