Coverage for /home/ubuntu/rolling-pin/python/rolling_pin/toml_etl.py: 100%

1from typing import Any, Type, TypeVar, Union # noqa: F401 

2 

3from copy import deepcopy 

4from pathlib import Path 

5import os 

6 

7from lunchbox.enforce import Enforce 

8import toml 

9 

10from rolling_pin.blob_etl import BlobETL 

11from toml.decoder import TomlDecodeError 

12 

13T = TypeVar('T', bound='TomlETL') 

14# ------------------------------------------------------------------------------ 

15 

16 

17class TomlETL: 

18 @classmethod 

19 def from_string(cls, text): 

20 # type: (Type[T], str) -> T 

21 ''' 

22 Creates a TomlETL instance from a given TOML string. 

23 

24 Args: 

25 text (str): TOML string. 

26 

27 Returns: 

28 TomlETL: TomlETL instance. 

29 ''' 

30 return cls(toml.loads(text)) 

31 

32 @classmethod 

33 def from_toml(cls, filepath): 

34 # type: (Type[T], Union[str, Path]) -> T 

35 ''' 

36 Creates a TomlETL instance from a given TOML file. 

37 

38 Args: 

39 filepath (str or Path): TOML file. 

40 

41 Returns: 

42 TomlETL: TomlETL instance. 

43 ''' 

44 return cls(toml.load(filepath)) 

45 

46 def __init__(self, data): 

47 # type: (dict[str, Any]) -> None 

48 ''' 

49 Creates a TomlETL instance from a given dictionary. 

50 

51 Args: 

52 data (dict): Dictionary. 

53 ''' 

54 self._data = data 

55 

56 def to_dict(self): 

57 # type: () -> dict 

58 ''' 

59 Converts instance to dictionary copy. 

60 

61 Returns: 

62 dict: Dictionary copy of instance. 

63 ''' 

64 return deepcopy(self._data) 

65 

66 def to_string(self): 

67 # type: () -> str 

68 ''' 

69 Converts instance to a TOML formatted string. 

70 

71 Returns: 

72 str: TOML string. 

73 ''' 

74 return toml.dumps( 

75 self._data, encoder=toml.TomlArraySeparatorEncoder(separator=',') 

76 ) 

77 

78 def write(self, filepath): 

79 # type: (Union[str, Path]) -> None 

80 ''' 

81 Writes instance to given TOML file. 

82 

83 Args: 

84 filepath (str or Path): Target filepath. 

85 ''' 

86 filepath = Path(filepath) 

87 os.makedirs(filepath.parent, exist_ok=True) 

88 with open(filepath, 'w') as f: 

89 toml.dump( 

90 self._data, 

91 f, 

92 encoder=toml.TomlArraySeparatorEncoder(separator=',') 

93 ) 

94 

95 def edit(self, patch): 

96 # type: (str) -> TomlETL 

97 ''' 

98 Apply edit to internal data given TOML patch. 

99 Patch is always of the form '[key]=[value]' and in TOML format. 

100 

101 Args: 

102 patch (str): TOML patch to be applied. 

103 

104 Raises: 

105 TOMLDecoderError: If patch cannot be decoded. 

106 EnforceError: If '=' not found in patch. 

107 

108 Returns: 

109 TomlETL: New TomlETL instance with edits. 

110 ''' 

111 msg = 'Edit patch must be a TOML parsable key value snippet with a "=" ' 

112 msg += 'character.' 

113 try: 

114 toml.loads(patch) 

115 except TomlDecodeError as e: 

116 msg += ' ' + e.msg 

117 raise TomlDecodeError(msg, e.doc, e.pos) 

118 Enforce('=', 'in', patch, message=msg) 

119 # ---------------------------------------------------------------------- 

120 

121 key, val = patch.split('=', maxsplit=1) 

122 val = toml.loads(f'x={val}')['x'] 

123 data = BlobETL(self._data, separator='.').to_flat_dict() 

124 data[key] = val 

125 data = BlobETL(data, separator='.').to_dict() 

126 return TomlETL(data) 

127 

128 def delete(self, regex): 

129 # type: (str) -> TomlETL 

130 ''' 

131 Returns portion of data whose keys fo not match a given regular expression. 

132 

133 Args: 

134 regex (str): Regular expression applied to keys. 

135 

136 Returns: 

137 TomlETL: New TomlETL instance. 

138 ''' 

139 data = BlobETL(self._data, separator='.') \ 

140 .query(regex, ignore_case=False, invert=True) \ 

141 .to_dict() 

142 return TomlETL(data) 

143 

144 def search(self, regex): 

145 # type: (str) -> TomlETL 

146 ''' 

147 Returns portion of data whose keys match a given regular expression. 

148 

149 Args: 

150 regex (str): Regular expression applied to keys. 

151 

152 Returns: 

153 TomlETL: New TomlETL instance. 

154 ''' 

155 data = BlobETL(self._data, separator='.') \ 

156 .query(regex, ignore_case=False) \ 

157 .to_dict() 

158 return TomlETL(data)