load.py (1030B)
1 from typing import Any, Dict, IO 2 from ..meta.schema import SchemaValue 3 4 import yaml 5 6 # PyYaml does not currently handle unique keys. 7 # https://github.com/yaml/pyyaml/issues/165#issuecomment-430074049 8 # In that issue, there are workarounds to it. 9 # https://gist.github.com/pypt/94d747fe5180851196eb?permalink_comment_id=4015118#gistcomment-4015118 10 11 class UniqueKeyLoader(yaml.SafeLoader): 12 def construct_mapping(self, node: yaml.MappingNode, deep: bool = False) -> Dict[Any, Any]: 13 mapping = set() 14 for key_node, value_node in node.value: 15 key = self.construct_object(key_node, deep=deep) # type: ignore 16 if key in mapping: 17 raise ValueError(f"Duplicate {key!r} key found in YAML.") 18 mapping.add(key) 19 return super().construct_mapping(node, deep) 20 21 def load_data_to_dict(f: IO[bytes]) -> Dict[str, Any]: 22 try: 23 raw_data = yaml.load(f, Loader=UniqueKeyLoader) 24 return SchemaValue.from_dict(raw_data) 25 except Exception as e: 26 raise e