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
« 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.
15"""Common functionality used across CSET."""
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
28import ruamel.yaml
30logger = logging.getLogger(__name__)
33class ArgumentError(ValueError):
34 """Provided arguments are not understood."""
37def sample_data_path(name):
38 """Return absolute path to sample data file.
40 Given the name of requested sample data resource, returns the full
41 path to the file.
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.
47 Parameters
48 ----------
49 name: str
50 The name of requested sample_data file.
52 Returns
53 -------
54 target: str
55 The full directory path to the requested file.
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)
73def parse_recipe(recipe_yaml: Path | str, variables: dict | None = None) -> dict:
74 """Parse a recipe into a python dictionary.
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.
85 Returns
86 -------
87 recipe: dict
88 The recipe as a python dictionary.
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.
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.")
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
117 logger.debug("Recipe before templating:\n%s", recipe)
118 check_recipe_has_steps(recipe)
120 if variables is not None:
121 logger.debug("Recipe variables: %s", variables)
122 recipe = template_variables(recipe, variables)
124 logger.debug("Recipe after templating:\n%s", recipe)
125 return recipe
128def check_recipe_has_steps(recipe: dict):
129 """Check a recipe has the minimum required steps.
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.
135 Parameters
136 ----------
137 recipe: dict
138 The recipe as a python dictionary.
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
160def slugify(s: str) -> str:
161 """Turn a string into a version that can be used everywhere.
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("_")
169def filename_slugify(s: str) -> str:
170 """Turn a string into a version that can be used in filenames.
172 The resultant string will only consist of a-z, 0-9.
173 """
174 return re.sub(r"[^a-z0-9\.]+", "", s.casefold()).strip("_")
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 {}
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.
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.
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.
205 Returns
206 -------
207 recipe_variables: dict
208 Dictionary keyed with the variable names.
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
245def template_variables(recipe: dict | list, variables: dict) -> dict:
246 """Insert variables into recipe.
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.
255 Returns
256 -------
257 recipe: dict
258 Filled recipe as a python dictionary.
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)
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
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
308################################################################################
309# Templating code taken from the simple_template package under the 0BSD licence.
310# Original at https://github.com/Fraetor/simple_template
311################################################################################
314class TemplateError(KeyError):
315 """Rendering a template failed due a placeholder without a value."""
318def render(template: str, /, **variables) -> str:
319 """Render the template with the provided variables.
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 }}`
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.
329 Parameters
330 ----------
331 template: str
332 Template to fill with variables.
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)`.
339 Returns
340 -------
341 rendered_template: str
342 Filled template.
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.
352 Examples
353 --------
354 >>> template = "<p>Hello {{myplaceholder}}!</p>"
355 >>> simple_template.render(template, myplaceholder="World")
356 "<p>Hello World!</p>"
357 """
359 def isidentifier(s: str):
360 return s.isidentifier()
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)
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)
375 for name in extract_placeholders():
376 template = substitute_placeholder(name)
377 return template
380def render_file(template_path: str, /, **variables) -> str:
381 """Render a template directly from a file.
383 Otherwise the same as `simple_template.render()`.
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)
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,)
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
406 def alphanum_key(s):
407 """Turn a string into a list of string and number chunks.
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
417 return sorted(iterable, key=alphanum_key, reverse=reverse)
420def combine_dicts(d1: dict, d2: dict) -> dict:
421 """Recursively combines two dictionaries.
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
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 }
446def sstrip(text):
447 """Dedent and strip text.
449 Parameters
450 ----------
451 text: str
452 The string to strip.
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()
468def is_increasing(sequence: list) -> bool:
469 """Determine the direction of an ordered sequence.
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]