Coverage for haystack/marshal/yaml.py: 95%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import Any 

6 

7import yaml 

8 

9 

10# Custom YAML safe loader that supports loading Python tuples 

11class YamlLoader(yaml.SafeLoader): 

12 def construct_python_tuple(self, node: yaml.SequenceNode) -> tuple: 

13 """Construct a Python tuple from the sequence.""" 

14 return tuple(self.construct_sequence(node)) 

15 

16 

17class YamlDumper(yaml.SafeDumper): 

18 def represent_tuple(self, data: tuple) -> yaml.SequenceNode: 

19 """Represent a Python tuple.""" 

20 return self.represent_sequence("tag:yaml.org,2002:python/tuple", data) 

21 

22 

23YamlDumper.add_representer(tuple, YamlDumper.represent_tuple) 

24YamlLoader.add_constructor("tag:yaml.org,2002:python/tuple", YamlLoader.construct_python_tuple) 

25 

26 

27class YamlMarshaller: 

28 def marshal(self, dict_: dict[str, Any]) -> str: 

29 """Return a YAML representation of the given dictionary.""" 

30 try: 

31 return yaml.dump(dict_, Dumper=YamlDumper) 

32 except yaml.representer.RepresenterError as e: 

33 raise TypeError( 

34 "Error dumping pipeline to YAML - Ensure that all pipeline components only serialize basic Python types" 

35 ) from e 

36 

37 def unmarshal(self, data_: str | bytes | bytearray) -> dict[str, Any]: 

38 """Return a dictionary from the given YAML data.""" 

39 try: 

40 return yaml.load(data_, Loader=YamlLoader) 

41 except yaml.constructor.ConstructorError as e: 

42 raise TypeError( 

43 "Error loading pipeline from YAML - Ensure that all pipeline " 

44 "components only serialize basic Python types" 

45 ) from e