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

1 import shlex

2 import typing

3 from dataclasses import dataclass, field

4 from pathlib import Path

5 from typing import Optional, Self, get_type_hints

6

7 import tomlkit

8

9

10@dataclass

11 class 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