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

156 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-07 15:12 +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, Sequence 

23from importlib.resources import files 

24from pathlib import Path 

25from textwrap import dedent 

26from typing import Any 

27 

28import ruamel.yaml 

29 

30logger = logging.getLogger(__name__) 

31 

32 

33class ArgumentError(ValueError): 

34 """Provided arguments are not understood.""" 

35 

36 

37def sample_data_path(name): 

38 """Return absolute path to sample data file. 

39 

40 Given the name of requested sample data resource, returns the full 

41 path to the file. 

42 

43 Note this function is only for locating files in the sample_data 

44 collection of files, used for generating documentation. It is not 

45 needed for general file access. 

46 

47 Parameters 

48 ---------- 

49 name: str 

50 The name of requested sample_data file. 

51 

52 Returns 

53 ------- 

54 target: str 

55 The full directory path to the requested file. 

56 

57 Raises 

58 ------ 

59 ValueError 

60 If the requested sample data is not found. 

61 """ 

62 target = files("CSET.sample_data").joinpath(name) 

63 if not target.is_file(): 

64 raise ValueError( 

65 f"Sample data file {name!r} not found.\n" 

66 "NB This function is only for locating files in the " 

67 "CSET sample_data collection. It is not needed or " 

68 "appropriate for general file access." 

69 ) 

70 return str(target) 

71 

72 

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

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

75 

76 Parameters 

77 ---------- 

78 recipe_yaml: Path | str 

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

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

81 read. 

82 variables: dict 

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

84 

85 Returns 

86 ------- 

87 recipe: dict 

88 The recipe as a python dictionary. 

89 

90 Raises 

91 ------ 

92 ValueError 

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

94 TypeError 

95 If recipe_yaml isn't a Path or string. 

96 KeyError 

97 If needed recipe variables are not supplied. 

98 

99 Examples 

100 -------- 

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

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

103 """ 

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

105 if isinstance(recipe_yaml, str): 

106 recipe_yaml = io.StringIO(recipe_yaml) 

107 elif not isinstance(recipe_yaml, Path): 

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

109 

110 # Parse the recipe YAML. 

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

112 try: 

113 recipe = yaml.load(recipe_yaml) 

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

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

116 

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

118 check_recipe_has_steps(recipe) 

119 

120 if variables is not None: 

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

122 recipe = template_variables(recipe, variables) 

123 

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

125 return recipe 

126 

127 

128def check_recipe_has_steps(recipe: dict): 

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

130 

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

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

133 reads the raw data. 

134 

135 Parameters 

136 ---------- 

137 recipe: dict 

138 The recipe as a python dictionary. 

139 

140 Raises 

141 ------ 

142 ValueError 

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

144 TypeError 

145 If recipe isn't a dict. 

146 KeyError 

147 If needed recipe variables are not supplied. 

148 """ 

149 if not isinstance(recipe, dict): 

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

151 if "steps" not in recipe: 

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

153 try: 

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

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

156 except TypeError as err: 

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

158 

159 

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

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

162 

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

164 underscores. 

165 """ 

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

167 

168 

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

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

171 

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

173 """ 

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

175 

176 

177def get_recipe_metadata() -> dict: 

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

179 try: 

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

181 return json.load(fp) 

182 except FileNotFoundError: 

183 meta = {} 

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

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

186 return {} 

187 

188 

189def parse_variable_options( 

190 arguments: Sequence[str], input_dir: str | Sequence[str] | None = None 

191) -> dict: 

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

193 

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

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

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

197 

198 Parameters 

199 ---------- 

200 arguments: list[str] 

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

202 input_dir: str | list[str], optional 

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

204 

205 Returns 

206 ------- 

207 recipe_variables: dict 

208 Dictionary keyed with the variable names. 

209 

210 Raises 

211 ------ 

212 ValueError 

213 If any arguments cannot be parsed. 

214 """ 

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

216 if input_dir is not None: 

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

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

219 recipe_variables = {} 

220 i = 0 

221 while i < len(arguments): 

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

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

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

225 try: 

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

227 value = arguments[i + 1] 

228 except IndexError as err: 

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

230 i += 1 

231 else: 

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

233 try: 

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

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

236 value = value[1:-1] 

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

238 # Capture the many possible exceptions from ast.literal_eval 

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

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

241 i += 1 

242 return recipe_variables 

243 

244 

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

246 """Insert variables into recipe. 

247 

248 Parameters 

249 ---------- 

250 recipe: dict | list 

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

252 variables: dict 

253 Dictionary of variables for the recipe. 

254 

255 Returns 

256 ------- 

257 recipe: dict 

258 Filled recipe as a python dictionary. 

259 

260 Raises 

261 ------ 

262 KeyError 

263 If needed recipe variables are not supplied. 

264 """ 

265 if isinstance(recipe, dict): 

266 index = recipe.keys() 

267 elif isinstance(recipe, list): 

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

269 index = range(len(recipe)) 

270 else: 

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

272 

273 for i in index: 

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

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

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

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

278 return recipe 

279 

280 

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

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

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

284 placeholder = f"${var_name}" 

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

286 # to keep the value type. 

287 if s == placeholder: 

288 # Specially handle Paths and lists of Paths. 

289 if isinstance(var_value, Path): 

290 var_value = str(var_value) 

291 if ( 

292 isinstance(var_value, list) 

293 and var_value 

294 and isinstance(var_value[0], Path) 

295 ): 

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

297 s = var_value 

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

299 # interpreting the new value. 

300 break 

301 else: 

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

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

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

305 return s 

306 

307 

308################################################################################ 

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

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

311################################################################################ 

312 

313 

314class TemplateError(KeyError): 

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

316 

317 

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

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

320 

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

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

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

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

325 

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

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

328 

329 Parameters 

330 ---------- 

331 template: str 

332 Template to fill with variables. 

333 

334 **variables: Any 

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

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

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

338 

339 Returns 

340 ------- 

341 rendered_template: str 

342 Filled template. 

343 

344 Raises 

345 ------ 

346 TemplateError 

347 Value not given for a placeholder in the template. 

348 TypeError 

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

350 string. 

351 

352 Examples 

353 -------- 

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

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

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

357 """ 

358 

359 def isidentifier(s: str): 

360 return s.isidentifier() 

361 

362 def extract_placeholders(): 

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

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

365 return filter(isidentifier, unique_names) 

366 

367 def substitute_placeholder(name): 

368 try: 

369 value = str(variables[name]) 

370 except KeyError as err: 

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

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

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

374 

375 for name in extract_placeholders(): 

376 template = substitute_placeholder(name) 

377 return template 

378 

379 

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

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

382 

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

384 

385 Examples 

386 -------- 

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

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

389 """ 

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

391 template = fp.read() 

392 return render(template, **variables) 

393 

394 

395def iter_maybe(thing) -> Iterable: 

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

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

398 return thing 

399 return (thing,) 

400 

401 

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

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

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

405 

406 def alphanum_key(s): 

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

408 

409 >>> alphanum_key("z23a") 

410 ["z", 23, "a"] 

411 """ 

412 try: 

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

414 except TypeError: 

415 return s 

416 

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

418 

419 

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

421 """Recursively combines two dictionaries. 

422 

423 Duplicate atoms favour the second dictionary. 

424 """ 

425 # Update existing keys. 

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

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

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

429 else: 

430 d1[key] = d2[key] 

431 # Add any new keys. 

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

433 d1[key] = d2[key] 

434 return d1 

435 

436 

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

438 """Recursively sort a dictionary.""" 

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

440 return { 

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

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

443 } 

444 

445 

446def sstrip(text): 

447 """Dedent and strip text. 

448 

449 Parameters 

450 ---------- 

451 text: str 

452 The string to strip. 

453 

454 Examples 

455 -------- 

456 >>> print(sstrip(''' 

457 ... foo 

458 ... bar 

459 ... baz 

460 ... ''')) 

461 foo 

462 bar 

463 baz 

464 """ 

465 return dedent(text).strip() 

466 

467 

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

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

470 

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

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

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

474 """ 

475 return sequence[0] < sequence[1]