Coverage for birdplan/yaml/pyaml.py: 100%

1# 

2# SPDX-License-Identifier: GPL-3.0-or-later 

3# 

4# Copyright (C) 2019-2024, AllWorldIT. 

5# 

6# This program is free software: you can redistribute it and/or modify 

7# it under the terms of the GNU General Public License as published by 

8# the Free Software Foundation, either version 3 of the License, or 

9# (at your option) any later version. 

10# 

11# This program is distributed in the hope that it will be useful, 

12# but WITHOUT ANY WARRANTY; without even the implied warranty of 

13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

14# GNU General Public License for more details. 

15# 

16# You should have received a copy of the GNU General Public License 

17# along with this program. If not, see <https://www.gnu.org/licenses/>. 

18 

19"""BirdPlan YAML handling.""" 

20 

21import io 

22import pathlib 

23from typing import Any, Optional, Tuple, Union 

24 

25import yaml as pyaml 

26from yaml import YAMLError 

27 

28from .base import YAMLBase 

29 

30__all__ = ["YAML", "YAMLError"] 

31 

32 

33class BirdPlanSafeLoader(pyaml.SafeLoader): # pylint: disable=too-many-ancestors 

34 """Safe YAML loader wtih some specific datatypes.""" 

35 

36 def construct_python_tuple(self, node: Any) -> Tuple[Any, ...]: 

37 """Tuple constructor.""" 

38 return tuple(self.construct_sequence(node)) 

39 

40 

41class BirdPlanSafeDumper(pyaml.SafeDumper): 

42 """Safe YAML dumper with some specific datatypes.""" 

43 

44 def represent_tuple(self, data: Tuple[Any, ...]) -> pyaml.nodes.SequenceNode: 

45 """Tuple representer.""" 

46 return self.represent_sequence("tag:yaml.org,2002:python/tuple", data, flow_style=True) 

47 

48 

49BirdPlanSafeLoader.add_constructor("tag:yaml.org,2002:python/tuple", BirdPlanSafeLoader.construct_python_tuple) 

50BirdPlanSafeDumper.add_representer(tuple, BirdPlanSafeDumper.represent_tuple) 

51 

52 

53class YAML(YAMLBase): 

54 """YAML class.""" 

55 

56 def __init__(self) -> None: 

57 """Initialize our parser YAML.""" 

58 

59 def load(self, yaml: Union[str, pathlib.Path, io.IOBase]) -> Any: 

60 """Load YAML string.""" 

61 

62 yaml_data: str = "" 

63 # Handle strings 

64 if isinstance(yaml, str): 

65 yaml_data = yaml 

66 # Handle path objects 

67 elif isinstance(yaml, pathlib.Path): 

68 with open(yaml, "r", encoding="UTF-8") as yaml_file: 

69 yaml_data = yaml_file.read() 

70 # Whats left over is file objects 

71 elif isinstance(yaml, io.IOBase): 

72 yaml_data = yaml.read() 

73 

74 return pyaml.load(yaml_data, BirdPlanSafeLoader) # nosec 

75 

76 def dump(self, data: Any, stream: Optional[Union[pathlib.Path, io.IOBase]] = None) -> Any: 

77 """Dump to YAML.""" 

78 

79 if stream: 

80 # Handle path objects 

81 if isinstance(stream, pathlib.Path): 

82 with open(stream, "w", encoding="UTF-8") as dump_file: 

83 return pyaml.dump(data, stream=dump_file, encoding="UTF-8", Dumper=BirdPlanSafeDumper) 

84 # Whats left over is file objects 

85 return pyaml.dump( 

86 data, 

87 stream=stream, 

88 encoding="UTF-8", 

89 Dumper=BirdPlanSafeDumper, 

90 ) 

91 

92 # Create a string IO object and dump the YAML data to it 

93 dumpstr = io.StringIO() # pragma: no cover 

94 with dumpstr: # pragma: no cover 

95 pyaml.dump(data, dumpstr, encoding="UTF-8", Dumper=BirdPlanSafeDumper) 

96 return dumpstr.getvalue()