Coverage for src/edwh_restic_plugin/forget.py: 87%

130 statements  

« prev     ^ index     » next       coverage.py v7.6.12, created at 2025-02-28 16:31 +0100

1import shlex 

2import typing 

3from dataclasses import dataclass, field 

4from pathlib import Path 

5from typing import Optional, Self, get_type_hints 

6 

7import tomlkit 

8 

9 

10@dataclass 

11class ResticForgetPolicy: 

12 """ 

13 Represents a policy for forgetting backups using restic, with various retention options. 

14 

15 Attributes: 

16 keep_last (Optional[int]): Number of latest snapshots to keep. 

17 keep_hourly (Optional[int]): Number of hourly snapshots to keep. 

18 keep_daily (Optional[int]): Number of daily snapshots to keep. 

19 keep_weekly (Optional[int]): Number of weekly snapshots to keep. 

20 keep_monthly (Optional[int]): Number of monthly snapshots to keep. 

21 keep_yearly (Optional[int]): Number of yearly snapshots to keep. 

22 keep_tag (list[str]): List of tags for which to apply the retention policy. 

23 keep_within (Optional[str]): Retention period within which to keep snapshots. 

24 keep_within_hourly (Optional[str]): Hourly retention period within which to keep snapshots. 

25 keep_within_daily (Optional[str]): Daily retention period within which to keep snapshots. 

26 keep_within_weekly (Optional[str]): Weekly retention period within which to keep snapshots. 

27 keep_within_monthly (Optional[str]): Monthly retention period within which to keep snapshots. 

28 keep_within_yearly (Optional[str]): Yearly retention period within which to keep snapshots. 

29 purge (bool): Whether to purge old snapshots. Default is True. 

30 """ 

31 

32 keep_last: Optional[int] = None 

33 keep_hourly: Optional[int] = None 

34 keep_daily: Optional[int] = None 

35 keep_weekly: Optional[int] = None 

36 keep_monthly: Optional[int] = None 

37 keep_yearly: Optional[int] = None 

38 keep_tag: list[str] = field(default_factory=list) 

39 keep_within: Optional[str] = None 

40 keep_within_hourly: Optional[str] = None 

41 keep_within_daily: Optional[str] = None 

42 keep_within_weekly: Optional[str] = None 

43 keep_within_monthly: Optional[str] = None 

44 keep_within_yearly: Optional[str] = None 

45 purge: bool = True 

46 

47 def to_string(self) -> str: 

48 """ 

49 Converts the policy into a string of restic command line arguments. 

50 

51 Returns: 

52 str: A string representing the restic command line arguments for this policy. 

53 """ 

54 args = [] 

55 for attr, value in vars(self).items(): 

56 if value is not None: 

57 option = "--" + attr.replace("_", "-") 

58 if isinstance(value, list): 

59 args.extend(f"{option} {shlex.quote(str(v))}" for v in value) 

60 elif isinstance(value, bool): 

61 args.append(option) 

62 else: 

63 args.append(f"{option} {shlex.quote(str(value))}") 

64 return " ".join(args) 

65 

66 @classmethod 

67 def from_string(cls, *args: str) -> Self: 

68 """ 

69 Creates a policy instance from a string of restic command line arguments. 

70 

71 Args: 

72 *args (str): A variable number of strings representing the restic command line arguments. 

73 

74 Returns: 

75 ResticForgetPolicy: An instance of ResticForgetPolicy created from the provided arguments. 

76 

77 Raises: 

78 ValueError: If a required argument is missing or incorrectly formatted. 

79 """ 

80 options = {} 

81 parsed_args = shlex.split(" ".join(args)) 

82 type_hints = get_type_hints(cls) 

83 

84 iterator = iter(parsed_args) 

85 for arg in iterator: 

86 if arg.startswith("--"): 

87 key = arg[2:].replace("-", "_") 

88 

89 # Handle boolean flags with --no- prefix 

90 if key.startswith("no_"): 

91 key = key[3:] 

92 attr_type = type_hints.get(key, bool) 

93 if attr_type is bool: 

94 options[key] = False 

95 continue 

96 

97 attr_type = type_hints.get(key, str) 

98 

99 # Handle boolean flags without value 

100 if attr_type is bool: 

101 options[key] = True 

102 continue 

103 

104 if "=" in arg: 

105 key, value = arg[2:].split("=", 1) 

106 key = key.replace("-", "_") 

107 else: 

108 try: 

109 value = next(iterator) 

110 except StopIteration: 

111 raise ValueError(f"Missing value for {arg}") 

112 

113 # Determine if the attribute type is an integer or optional integer 

114 if typing.get_origin(attr_type) is typing.Union: 

115 if int in typing.get_args(attr_type): 

116 options[key] = int(value) 

117 elif attr_type is int: 

118 options[key] = int(value) 

119 elif attr_type is list: 

120 options.setdefault(key, []).append(value) 

121 else: 

122 options[key] = value 

123 

124 # Set default values for any missing boolean attributes 

125 for key, attr_type in type_hints.items(): 

126 if attr_type is bool and key not in options: 

127 options[key] = False 

128 

129 return cls(**options) 

130 

131 @classmethod 

132 def from_toml_file(cls, subkey: str, toml_path: Optional[str | Path] = None) -> Optional[Self]: 

133 """ 

134 Creates a policy instance from a TOML file. 

135 

136 Args: 

137 subkey (str): The key under which the policy is stored in the TOML file. 

138 toml_path (Optional[str | Path]): The path to the TOML configuration file. Defaults to '.toml' in the current directory. 

139 

140 Returns: 

141 ResticForgetPolicy: An instance of ResticForgetPolicy created from the provided TOML file and key, or None if not found. 

142 

143 Raises: 

144 FileNotFoundError: If the specified TOML file does not exist. 

145 KeyError: If the specified subkey is not found in the TOML file. 

146 """ 

147 # Set default toml_path if not provided 

148 if toml_path is None: 

149 toml_path = Path.cwd() / ".toml" 

150 else: 

151 toml_path = Path(toml_path) 

152 

153 try: 

154 data = tomlkit.parse(toml_path.read_text()) 

155 forget = data["restic"]["forget"] 

156 except (KeyError, OSError): 

157 return None 

158 

159 if not (section := (forget.get(subkey) or forget.get("default"))): 

160 return None 

161 

162 policy_dict = {} 

163 for key, value in section.items(): 

164 # Convert keys to snake_case 

165 snake_key = key.replace("-", "_") 

166 if snake_key == "purge": 

167 policy_dict[snake_key] = bool(value) 

168 else: 

169 try: 

170 value = int(value) 

171 except (ValueError, TypeError): 

172 pass # If it's not a valid integer, leave it as is 

173 policy_dict[snake_key] = value 

174 

175 return cls(**policy_dict) 

176 

177 def to_toml(self, toml_path: str | Path, subkey: str) -> None: 

178 """ 

179 Writes the policy to a TOML file under the specified subkey, replacing it if it exists. 

180 

181 Args: 

182 toml_path (str | Path): The path to the TOML configuration file. 

183 subkey (str): The key under which the policy should be stored in the TOML file. 

184 """ 

185 toml_path = Path(toml_path) 

186 

187 # Read existing TOML data or create new if not exists 

188 if toml_path.exists(): 

189 data = tomlkit.parse(toml_path.read_text()) 

190 else: 

191 data = tomlkit.document() 

192 

193 # Ensure 'restic' and 'forget' sections exist 

194 if "restic" not in data: 

195 data["restic"] = tomlkit.table() 

196 if "forget" not in data["restic"]: 

197 data["restic"]["forget"] = tomlkit.table() 

198 

199 # Prepare the policy data to write 

200 policy_data = {} 

201 for attr, value in vars(self).items(): 

202 if value is not None: 

203 key = attr.replace("_", "-") 

204 policy_data[key] = value 

205 

206 # Update the subkey 

207 data["restic"]["forget"][subkey] = policy_data 

208 

209 # Write back to the TOML file 

210 toml_path.write_text(tomlkit.dumps(data)) 

211 

212 @classmethod 

213 def get_or_copy_policy(cls, subkey: str, toml_path: Optional[str | Path] = None, 

214 default_toml_path: Optional[str | Path] = None) -> Optional[Self]: 

215 """ 

216 Retrieves a policy from the TOML file or copies it from the default TOML file if not present. 

217 

218 Args: 

219 subkey (str): The key under which the policy is stored in the TOML file. 

220 toml_path (Optional[str | Path]): The path to the TOML configuration file. Defaults to '.toml' in the current directory. 

221 default_toml_path (Optional[str | Path]): The path to the default TOML configuration file. Defaults to 'default.toml' in the same directory as toml_path. 

222 

223 Returns: 

224 ResticForgetPolicy: An instance of ResticForgetPolicy created from the provided TOML file and key, 

225 or copied from the default TOML file, or None if not found. 

226 """ 

227 # Ensure toml_path is a Path object 

228 if toml_path is None: 

229 toml_path = Path.cwd() / ".toml" 

230 else: 

231 toml_path = Path(toml_path) 

232 

233 # Set default_toml_path if not provided 

234 if default_toml_path is None: 

235 default_toml_path = toml_path.parent / "default.toml" 

236 else: 

237 default_toml_path = Path(default_toml_path) 

238 

239 # Try to get the policy from the main TOML file 

240 policy = cls.from_toml_file(subkey, toml_path) 

241 if policy: 

242 return policy 

243 

244 # Try to get the policy from the default TOML file 

245 policy = cls.from_toml_file(subkey, default_toml_path) 

246 if policy: 

247 # Copy the policy to the main TOML file 

248 policy.to_toml(toml_path, subkey) 

249 return policy 

250 

251 # Try to get the default policy from the default TOML file 

252 policy = cls.from_toml_file("default", default_toml_path) 

253 if policy: 

254 # Copy the default policy to the main TOML file under the subkey 

255 policy.to_toml(toml_path, subkey) 

256 return policy 

257 

258 return None 

259