Coverage for src/CSET/_common.py: 100%

150 statements  

« prev     ^ index     » next       coverage.py v7.15.3, created at 2026-08-04 08:32 +0000

1# © Crown copyright, Met Office (2022-2024) and CSET contributors. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Common functionality used across CSET.""" 

16 

17import ast 

18import io 

19import json 

20import logging 

21import re 

22from collections.abc import Iterable 

23from pathlib import Path 

24from textwrap import dedent 

25from typing import Any 

26 

27import ruamel.yaml 

28 

29logger = logging.getLogger(__name__) 

30 

31 

32class ArgumentError(ValueError): 

33 """Provided arguments are not understood.""" 

34 

35 

36def parse_recipe(recipe_yaml: Path | str, variables: dict | None = None) -> dict: 

37 """Parse a recipe into a python dictionary. 

38 

39 Parameters 

40 ---------- 

41 recipe_yaml: Path | str 

42 Path to a file containing, or a string of, a recipe's YAML describing 

43 the operators that need running. If a Path is provided it is opened and 

44 read. 

45 variables: dict 

46 Dictionary of recipe variables. If None templating is not attempted. 

47 

48 Returns 

49 ------- 

50 recipe: dict 

51 The recipe as a python dictionary. 

52 

53 Raises 

54 ------ 

55 ValueError 

56 If the recipe is invalid. E.g. invalid YAML, missing any steps, etc. 

57 TypeError 

58 If recipe_yaml isn't a Path or string. 

59 KeyError 

60 If needed recipe variables are not supplied. 

61 

62 Examples 

63 -------- 

64 >>> CSET._common.parse_recipe(Path("myrecipe.yaml")) 

65 {'steps': [{'operator': 'misc.noop'}]} 

66 """ 

67 # Ensure recipe_yaml is something the YAML parser can read. 

68 if isinstance(recipe_yaml, str): 

69 recipe_yaml = io.StringIO(recipe_yaml) 

70 elif not isinstance(recipe_yaml, Path): 

71 raise TypeError("recipe_yaml must be a str or Path.") 

72 

73 # Parse the recipe YAML. 

74 with ruamel.yaml.YAML(typ="safe", pure=True) as yaml: 

75 try: 

76 recipe = yaml.load(recipe_yaml) 

77 except ruamel.yaml.parser.ParserError as err: 

78 raise ValueError("ParserError: Invalid YAML") from err 

79 

80 logger.debug("Recipe before templating:\n%s", recipe) 

81 check_recipe_has_steps(recipe) 

82 

83 if variables is not None: 

84 logger.debug("Recipe variables: %s", variables) 

85 recipe = template_variables(recipe, variables) 

86 

87 logger.debug("Recipe after templating:\n%s", recipe) 

88 return recipe 

89 

90 

91def check_recipe_has_steps(recipe: dict): 

92 """Check a recipe has the minimum required steps. 

93 

94 Checking that the recipe actually has some steps, and providing helpful 

95 error messages otherwise. We must have at least a steps step, as that 

96 reads the raw data. 

97 

98 Parameters 

99 ---------- 

100 recipe: dict 

101 The recipe as a python dictionary. 

102 

103 Raises 

104 ------ 

105 ValueError 

106 If the recipe is invalid. E.g. invalid YAML, missing any steps, etc. 

107 TypeError 

108 If recipe isn't a dict. 

109 KeyError 

110 If needed recipe variables are not supplied. 

111 """ 

112 if not isinstance(recipe, dict): 

113 raise TypeError("Recipe must contain a mapping.") 

114 if "steps" not in recipe: 

115 raise ValueError("Recipe must contain a 'steps' key.") 

116 try: 

117 if len(recipe["steps"]) < 1: 

118 raise ValueError("Recipe must have at least 1 step.") 

119 except TypeError as err: 

120 raise ValueError("'steps' key must contain a sequence of steps.") from err 

121 

122 

123def slugify(s: str) -> str: 

124 """Turn a string into a version that can be used everywhere. 

125 

126 The resultant string will only consist of a-z, 0-9, dots, dashes, and 

127 underscores. 

128 """ 

129 return re.sub(r"[^a-z0-9\._-]+", "_", s.casefold()).strip("_") 

130 

131 

132def filename_slugify(s: str) -> str: 

133 """Turn a string into a version that can be used in filenames. 

134 

135 The resultant string will only consist of a-z, 0-9. 

136 """ 

137 return re.sub(r"[^a-z0-9\.]+", "", s.casefold()).strip("_") 

138 

139 

140def get_recipe_metadata() -> dict: 

141 """Get the metadata of the running recipe.""" 

142 try: 

143 with open("meta.json", "rt", encoding="UTF-8") as fp: 

144 return json.load(fp) 

145 except FileNotFoundError: 

146 meta = {} 

147 with open("meta.json", "wt", encoding="UTF-8") as fp: 

148 json.dump(meta, fp, indent=2) 

149 return {} 

150 

151 

152def parse_variable_options( 

153 arguments: list[str], input_dir: str | list[str] | None = None 

154) -> dict: 

155 """Parse a list of arguments into a dictionary of variables. 

156 

157 The variable name arguments start with two hyphen-minus (`--`), consisting 

158 of only capital letters (`A`-`Z`) and underscores (`_`). While the variable 

159 name is restricted, the value of the variable can be any string. 

160 

161 Parameters 

162 ---------- 

163 arguments: list[str] 

164 List of arguments, e.g: `["--LEVEL", "2", "--STASH=m01s01i001"]` 

165 input_dir: str | list[str], optional 

166 List of input directories to add into the returned variables. 

167 

168 Returns 

169 ------- 

170 recipe_variables: dict 

171 Dictionary keyed with the variable names. 

172 

173 Raises 

174 ------ 

175 ValueError 

176 If any arguments cannot be parsed. 

177 """ 

178 # Convert --input_dir=... to INPUT_PATHS recipe variable. 

179 if input_dir is not None: 

180 abs_paths = [str(Path(p).absolute()) for p in iter_maybe(input_dir)] 

181 arguments.append(f"--INPUT_PATHS={abs_paths}") 

182 recipe_variables = {} 

183 i = 0 

184 while i < len(arguments): 

185 if re.fullmatch(r"--[A-Z_]+=.*", arguments[i]): 

186 key, value = arguments[i].split("=", 1) 

187 elif re.fullmatch(r"--[A-Z_]+", arguments[i]): 

188 try: 

189 key = arguments[i].strip("-") 

190 value = arguments[i + 1] 

191 except IndexError as err: 

192 raise ArgumentError(f"No value for variable {arguments[i]}") from err 

193 i += 1 

194 else: 

195 raise ArgumentError(f"Unknown argument: {arguments[i]}") 

196 try: 

197 # Remove quotes from arguments, in case left in CSET_ADDOPTS. 

198 if re.fullmatch(r"""["'].+["']""", value): 

199 value = value[1:-1] 

200 recipe_variables[key.strip("-")] = ast.literal_eval(value) 

201 # Capture the many possible exceptions from ast.literal_eval 

202 except (ValueError, TypeError, SyntaxError, MemoryError, RecursionError): 

203 recipe_variables[key.strip("-")] = value 

204 i += 1 

205 return recipe_variables 

206 

207 

208def template_variables(recipe: dict | list, variables: dict) -> dict: 

209 """Insert variables into recipe. 

210 

211 Parameters 

212 ---------- 

213 recipe: dict | list 

214 The recipe as a python dictionary. It is updated in-place. 

215 variables: dict 

216 Dictionary of variables for the recipe. 

217 

218 Returns 

219 ------- 

220 recipe: dict 

221 Filled recipe as a python dictionary. 

222 

223 Raises 

224 ------ 

225 KeyError 

226 If needed recipe variables are not supplied. 

227 """ 

228 if isinstance(recipe, dict): 

229 index = recipe.keys() 

230 elif isinstance(recipe, list): 

231 # We have to handle lists for when we have one inside a recipe. 

232 index = range(len(recipe)) 

233 else: 

234 raise TypeError("recipe must be a dict or list.", recipe) 

235 

236 for i in index: 

237 if isinstance(recipe[i], (dict, list)): 

238 recipe[i] = template_variables(recipe[i], variables) 

239 elif isinstance(recipe[i], str): 

240 recipe[i] = replace_template_variable(recipe[i], variables) 

241 return recipe 

242 

243 

244def replace_template_variable(s: str, variables: dict[str, Any]): 

245 """Fill all variable placeholders in the string.""" 

246 for var_name, var_value in variables.items(): 

247 placeholder = f"${var_name}" 

248 # If the value is just the placeholder we directly overwrite it 

249 # to keep the value type. 

250 if s == placeholder: 

251 # Specially handle Paths and lists of Paths. 

252 if isinstance(var_value, Path): 

253 var_value = str(var_value) 

254 if ( 

255 isinstance(var_value, list) 

256 and var_value 

257 and isinstance(var_value[0], Path) 

258 ): 

259 var_value = [str(p) for p in var_value] 

260 s = var_value 

261 # We have replaced the whole string, so stop here to avoid 

262 # interpreting the new value. 

263 break 

264 else: 

265 s = s.replace(placeholder, str(var_value)) 

266 if isinstance(s, str) and re.match(r"^.*\$[A-Z_].*", s): 

267 raise KeyError("Variable without a value.", s) 

268 return s 

269 

270 

271################################################################################ 

272# Templating code taken from the simple_template package under the 0BSD licence. 

273# Original at https://github.com/Fraetor/simple_template 

274################################################################################ 

275 

276 

277class TemplateError(KeyError): 

278 """Rendering a template failed due a placeholder without a value.""" 

279 

280 

281def render(template: str, /, **variables) -> str: 

282 """Render the template with the provided variables. 

283 

284 The template should contain placeholders that will be replaced. These 

285 placeholders consist of the placeholder name within double curly braces. The 

286 name of the placeholder should be a valid python identifier. Whitespace 

287 between the braces and the name is ignored. E.g.: `{{ placeholder_name }}` 

288 

289 An exception will be raised if there are placeholders without corresponding 

290 values. It is acceptable to provide unused values; they will be ignored. 

291 

292 Parameters 

293 ---------- 

294 template: str 

295 Template to fill with variables. 

296 

297 **variables: Any 

298 Keyword arguments for the placeholder values. The argument name should 

299 be the same as the placeholder's name. You can unpack a dictionary of 

300 value with `render(template, **my_dict)`. 

301 

302 Returns 

303 ------- 

304 rendered_template: str 

305 Filled template. 

306 

307 Raises 

308 ------ 

309 TemplateError 

310 Value not given for a placeholder in the template. 

311 TypeError 

312 If the template is not a string, or a variable cannot be casted to a 

313 string. 

314 

315 Examples 

316 -------- 

317 >>> template = "<p>Hello {{myplaceholder}}!</p>" 

318 >>> simple_template.render(template, myplaceholder="World") 

319 "<p>Hello World!</p>" 

320 """ 

321 

322 def isidentifier(s: str): 

323 return s.isidentifier() 

324 

325 def extract_placeholders(): 

326 matches = re.finditer(r"{{\s*([^}]+)\s*}}", template) 

327 unique_names = {match.group(1) for match in matches} 

328 return filter(isidentifier, unique_names) 

329 

330 def substitute_placeholder(name): 

331 try: 

332 value = str(variables[name]) 

333 except KeyError as err: 

334 raise TemplateError("Placeholder missing value", name) from err 

335 pattern = r"{{\s*%s\s*}}" % re.escape(name) # noqa: UP031 Braces get ugly within f-string. 

336 return re.sub(pattern, value, template) 

337 

338 for name in extract_placeholders(): 

339 template = substitute_placeholder(name) 

340 return template 

341 

342 

343def render_file(template_path: str, /, **variables) -> str: 

344 """Render a template directly from a file. 

345 

346 Otherwise the same as `simple_template.render()`. 

347 

348 Examples 

349 -------- 

350 >>> simple_template.render_file("/path/to/template.html", myplaceholder="World") 

351 "<p>Hello World!</p>" 

352 """ 

353 with open(template_path, "rt", encoding="UTF-8") as fp: 

354 template = fp.read() 

355 return render(template, **variables) 

356 

357 

358def iter_maybe(thing) -> Iterable: 

359 """Ensure thing is Iterable. Strings count as atoms.""" 

360 if isinstance(thing, Iterable) and not isinstance(thing, str): 

361 return thing 

362 return (thing,) 

363 

364 

365def human_sorted(iterable: Iterable, reverse: bool = False) -> list: 

366 """Sort such numbers within strings are sorted correctly.""" 

367 # Adapted from https://nedbatchelder.com/blog/200712/human_sorting.html 

368 

369 def alphanum_key(s): 

370 """Turn a string into a list of string and number chunks. 

371 

372 >>> alphanum_key("z23a") 

373 ["z", 23, "a"] 

374 """ 

375 try: 

376 return [int(c) if c.isdecimal() else c for c in re.split(r"(\d+)", s)] 

377 except TypeError: 

378 return s 

379 

380 return sorted(iterable, key=alphanum_key, reverse=reverse) 

381 

382 

383def combine_dicts(d1: dict, d2: dict) -> dict: 

384 """Recursively combines two dictionaries. 

385 

386 Duplicate atoms favour the second dictionary. 

387 """ 

388 # Update existing keys. 

389 for key in d1.keys() & d2.keys(): 

390 if isinstance(d1[key], dict): 

391 d1[key] = combine_dicts(d1[key], d2[key]) 

392 else: 

393 d1[key] = d2[key] 

394 # Add any new keys. 

395 for key in d2.keys() - d1.keys(): 

396 d1[key] = d2[key] 

397 return d1 

398 

399 

400def sort_dict(d: dict) -> dict: 

401 """Recursively sort a dictionary.""" 

402 # Thank you to https://stackoverflow.com/a/47882384 

403 return { 

404 k: sort_dict(v) if isinstance(v, dict) else v 

405 for k, v in human_sorted(d.items()) 

406 } 

407 

408 

409def sstrip(text): 

410 """Dedent and strip text. 

411 

412 Parameters 

413 ---------- 

414 text: str 

415 The string to strip. 

416 

417 Examples 

418 -------- 

419 >>> print(sstrip(''' 

420 ... foo 

421 ... bar 

422 ... baz 

423 ... ''')) 

424 foo 

425 bar 

426 baz 

427 """ 

428 return dedent(text).strip() 

429 

430 

431def is_increasing(sequence: list) -> bool: 

432 """Determine the direction of an ordered sequence. 

433 

434 Returns a boolean indicating that the values of a sequence are 

435 increasing. The sequence should already be monotonic, with no 

436 duplicate values. An iris DimCoord's points fulfils this criteria. 

437 """ 

438 return sequence[0] < sequence[1]