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
« 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.
15"""Common functionality used across CSET."""
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
27import ruamel.yaml
29logger = logging.getLogger(__name__)
32class ArgumentError(ValueError):
33 """Provided arguments are not understood."""
36def parse_recipe(recipe_yaml: Path | str, variables: dict | None = None) -> dict:
37 """Parse a recipe into a python dictionary.
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.
48 Returns
49 -------
50 recipe: dict
51 The recipe as a python dictionary.
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.
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.")
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
80 logger.debug("Recipe before templating:\n%s", recipe)
81 check_recipe_has_steps(recipe)
83 if variables is not None:
84 logger.debug("Recipe variables: %s", variables)
85 recipe = template_variables(recipe, variables)
87 logger.debug("Recipe after templating:\n%s", recipe)
88 return recipe
91def check_recipe_has_steps(recipe: dict):
92 """Check a recipe has the minimum required steps.
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.
98 Parameters
99 ----------
100 recipe: dict
101 The recipe as a python dictionary.
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
123def slugify(s: str) -> str:
124 """Turn a string into a version that can be used everywhere.
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("_")
132def filename_slugify(s: str) -> str:
133 """Turn a string into a version that can be used in filenames.
135 The resultant string will only consist of a-z, 0-9.
136 """
137 return re.sub(r"[^a-z0-9\.]+", "", s.casefold()).strip("_")
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 {}
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.
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.
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.
168 Returns
169 -------
170 recipe_variables: dict
171 Dictionary keyed with the variable names.
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
208def template_variables(recipe: dict | list, variables: dict) -> dict:
209 """Insert variables into recipe.
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.
218 Returns
219 -------
220 recipe: dict
221 Filled recipe as a python dictionary.
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)
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
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
271################################################################################
272# Templating code taken from the simple_template package under the 0BSD licence.
273# Original at https://github.com/Fraetor/simple_template
274################################################################################
277class TemplateError(KeyError):
278 """Rendering a template failed due a placeholder without a value."""
281def render(template: str, /, **variables) -> str:
282 """Render the template with the provided variables.
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 }}`
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.
292 Parameters
293 ----------
294 template: str
295 Template to fill with variables.
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)`.
302 Returns
303 -------
304 rendered_template: str
305 Filled template.
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.
315 Examples
316 --------
317 >>> template = "<p>Hello {{myplaceholder}}!</p>"
318 >>> simple_template.render(template, myplaceholder="World")
319 "<p>Hello World!</p>"
320 """
322 def isidentifier(s: str):
323 return s.isidentifier()
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)
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)
338 for name in extract_placeholders():
339 template = substitute_placeholder(name)
340 return template
343def render_file(template_path: str, /, **variables) -> str:
344 """Render a template directly from a file.
346 Otherwise the same as `simple_template.render()`.
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)
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,)
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
369 def alphanum_key(s):
370 """Turn a string into a list of string and number chunks.
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
380 return sorted(iterable, key=alphanum_key, reverse=reverse)
383def combine_dicts(d1: dict, d2: dict) -> dict:
384 """Recursively combines two dictionaries.
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
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 }
409def sstrip(text):
410 """Dedent and strip text.
412 Parameters
413 ----------
414 text: str
415 The string to strip.
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()
431def is_increasing(sequence: list) -> bool:
432 """Determine the direction of an ordered sequence.
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]