Coverage for src/CSET/recipes/__init__.py: 100%

1# © Crown copyright, Met Office (2022-2025) 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"""Operations on recipes.""" 

16 

17import hashlib 

18import importlib.resources 

19import logging 

20from collections.abc import Iterator 

21from io import StringIO 

22from pathlib import Path 

23from typing import Any 

24 

25from ruamel.yaml import YAML 

26 

27from CSET._common import parse_recipe, slugify 

28from CSET.cset_workflow.lib.python.jinja_utils import get_models as get_models 

29 

30logger = logging.getLogger(__name__) 

31 

32 

33def _recipe_files_in_tree( 

34 recipe_name: str | None = None, input_dir: Path | None = None 

35) -> Iterator[Path]: 

36 """Yield recipe file Paths matching the recipe name.""" 

37 if input_dir is None: 

38 input_dir = importlib.resources.files() 

39 for file in input_dir.iterdir(): 

40 logger.debug("Testing %s", file) 

41 if ( 

42 (recipe_name is None or recipe_name == file.name) 

43 and file.is_file() 

44 and file.suffix == ".yaml" 

45 ): 

46 yield file 

47 elif file.is_dir() and file.name[0] != "_": # Excludes __pycache__ 

48 yield from _recipe_files_in_tree(recipe_name, file) 

49 

50 

51def _get_recipe_file(recipe_name: str, input_dir: Path | None = None) -> Path: 

52 """Return a Path to the recipe file.""" 

53 if input_dir is None: 

54 input_dir = importlib.resources.files() 

55 file = input_dir / recipe_name 

56 logger.debug("Getting recipe: %s", file) 

57 if not file.is_file(): 

58 raise FileNotFoundError("Recipe file does not exist.", recipe_name) 

59 return file 

60 

61 

62def unpack_recipe(recipe_dir: Path, recipe_name: str) -> None: 

63 """ 

64 Unpacks recipes files into a directory, creating it if it doesn't exist. 

65 

66 Parameters 

67 ---------- 

68 recipe_dir: Path 

69 Path to a directory into which to unpack the recipe files. 

70 recipe_name: str 

71 Name of recipe to unpack. 

72 

73 Raises 

74 ------ 

75 FileExistsError 

76 If recipe_dir already exists, and is not a directory. 

77 

78 OSError 

79 If recipe_dir cannot be created, such as insufficient permissions, or 

80 lack of space. 

81 """ 

82 recipe_dir.mkdir(parents=True, exist_ok=True) 

83 output_file = recipe_dir / recipe_name 

84 logger.debug("Saving recipe to %s", output_file) 

85 if output_file.exists(): 

86 logger.debug("%s already exists in target directory, skipping.", recipe_name) 

87 return 

88 logger.info("Unpacking %s to %s", recipe_name, output_file) 

89 try: 

90 file = _get_recipe_file(next(_recipe_files_in_tree(recipe_name))) 

91 except StopIteration as err: 

92 raise FileNotFoundError("Recipe not found.") from err 

93 output_file.write_bytes(file.read_bytes()) 

94 

95 

96def list_available_recipes() -> None: 

97 """List available recipes to stdout.""" 

98 print("Available recipes:") 

99 for file in _recipe_files_in_tree(): 

100 print(f"\t{file.name}") 

101 

102 

103def detail_recipe(recipe_name: str) -> None: 

104 """Detail the recipe to stdout. 

105 

106 If multiple recipes match the given name they will all be displayed. 

107 

108 Parameters 

109 ---------- 

110 recipe_name: str 

111 Partial match for the recipe name. 

112 """ 

113 for file in _recipe_files_in_tree(recipe_name): 

114 with YAML(typ="safe", pure=True) as yaml: 

115 recipe = yaml.load(file) 

116 print(f"\n\t{file.name}\n\t{''.join('─' * len(file.name))}\n") 

117 print(recipe.get("description")) 

118 

119 

120class RawRecipe: 

121 """A recipe to be parbaked. 

122 

123 Parameters 

124 ---------- 

125 recipe: str 

126 Name of the recipe file. 

127 model_ids: int | list[int] 

128 Model IDs to set the input paths for. Matches the corresponding workflow 

129 model IDs. 

130 variables: dict[str, Any] aggregation: bool 

131 Recipe variables to be inserted into $VAR placeholders in the recipe. 

132 aggregation: bool 

133 Whether this is an aggregation recipe or just a single case. 

134 

135 Returns 

136 ------- 

137 RawRecipe 

138 """ 

139 

140 recipe: str 

141 model_ids: list[int] 

142 variables: dict[str, Any] 

143 aggregation: bool 

144 

145 def __init__( 

146 self, 

147 recipe: str, 

148 model_ids: int | list[int], 

149 variables: dict[str, Any], 

150 aggregation: bool, 

151 ) -> None: 

152 self.recipe = recipe 

153 self.model_ids = model_ids if isinstance(model_ids, list) else [model_ids] 

154 self.variables = variables 

155 self.aggregation = aggregation 

156 

157 def __str__(self) -> str: 

158 """Return str(self). 

159 

160 Examples 

161 -------- 

162 >>> print(raw_recipe) 

163 generic_surface_spatial_plot_sequence.yaml (model 1) 

164 VARNAME air_temperature 

165 MODEL_NAME Model A 

166 METHOD SEQ 

167 SUBAREA_TYPE None 

168 SUBAREA_EXTENT None 

169 """ 

170 recipe = self.recipe if self.recipe else "<unknown>" 

171 plural = "s" if len(self.model_ids) > 1 else "" 

172 ids = " ".join(str(m) for m in self.model_ids) 

173 aggregation = ", Aggregation" if self.aggregation else "" 

174 pad = max([0] + [len(k) for k in self.variables.keys()]) 

175 variables = "".join(f"\n\t{k:<{pad}} {v}" for k, v in self.variables.items()) 

176 return f"{recipe} (model{plural} {ids}{aggregation}){variables}" 

177 

178 def __eq__(self, value: object) -> bool: 

179 """Return self==value.""" 

180 if isinstance(value, self.__class__): 

181 return ( 

182 self.recipe == value.recipe 

183 and self.model_ids == value.model_ids 

184 and self.variables == value.variables 

185 and self.aggregation == value.aggregation 

186 ) 

187 return NotImplemented 

188 

189 def parbake(self, ROSE_DATAC: Path, SHARE_DIR: Path) -> None: 

190 """Pre-process recipe to bake in all variables. 

191 

192 Parameters 

193 ---------- 

194 ROSE_DATAC: Path 

195 Workflow shared per-cycle data location. 

196 SHARE_DIR: Path 

197 Workflow shared data location. 

198 """ 

199 # Ready recipe file to disk. 

200 unpack_recipe(Path.cwd(), self.recipe) 

201 

202 # Collect configuration from environment. 

203 if self.aggregation: 

204 # Construct the location for the recipe. 

205 recipe_dir = ROSE_DATAC / "aggregation_recipes" 

206 # Construct the input data directories for the cycle. 

207 data_dirs = [ 

208 SHARE_DIR / f"cycle/*/data/{model_id}" for model_id in self.model_ids 

209 ] 

210 else: 

211 recipe_dir = ROSE_DATAC / "recipes" 

212 data_dirs = [ROSE_DATAC / f"data/{model_id}" for model_id in self.model_ids] 

213 

214 # Ensure recipe dir exists. 

215 recipe_dir.mkdir(parents=True, exist_ok=True) 

216 

217 # Add input paths to recipe variables. 

218 self.variables["INPUT_PATHS"] = data_dirs 

219 

220 # Parbake this recipe. 

221 recipe = parse_recipe(Path(self.recipe), self.variables) 

222 

223 # Add variables as extra metadata to filter on. 

224 for key, value in self.variables.items(): 

225 # Don't overwrite existing keys. 

226 if key != "INPUT_PATHS" and key not in recipe: 

227 recipe[key] = value 

228 

229 # Serialise into memory, as we use the serialised value twice. 

230 with StringIO() as s: 

231 with YAML(pure=True, output=s) as yaml: 

232 yaml.dump(recipe) 

233 serialised_recipe = s.getvalue().encode() 

234 # Include shortened hash in filename to avoid collisions between recipes 

235 # with the same title. 

236 digest = hashlib.sha256(serialised_recipe).hexdigest() 

237 output_filename = recipe_dir / f"{slugify(recipe['title'])}_{digest[:12]}.yaml" 

238 # Save into recipe_dir. 

239 with open(output_filename, "wb") as fp: 

240 fp.write(serialised_recipe) 

241 

242 

243class Config: 

244 """Namespace for easy access to configuration values. 

245 

246 A namespace for easy access to configuration values (via config.variable), 

247 where undefined attributes return an empty list. An empty list evaluates to 

248 False in boolean contexts and can be safely iterated over, so it acts as an 

249 effective unset value. 

250 

251 Parameters 

252 ---------- 

253 config: dict 

254 Configuration key-value pairs. 

255 

256 Example 

257 ------- 

258 >>> conf = Config({"key": "value"}) 

259 >>> conf.key 

260 'value' 

261 >>> conf.missing 

262 [] 

263 """ 

264 

265 d: dict 

266 

267 def __init__(self, config: dict) -> None: 

268 self.d = config 

269 

270 def __getattr__(self, name: str): 

271 """Return an empty list for missing names.""" 

272 return self.d.get(name, []) 

273 

274 def asdict(self) -> dict: 

275 """Return config as a dictionary.""" 

276 return self.d 

277 

278 

279def load_recipes(variables: dict[str, Any]) -> Iterator[RawRecipe]: 

280 """Load recipes enabled by configuration. 

281 

282 Recipes are loaded using all loaders (python modules) in CSET.loaders. Each 

283 of these loaders must define a function with the signature `load(conf: dict) 

284 -> Iterator[RawRecipe]`, which will be called with `variables`. 

285 

286 A minimal example can be found in `CSET.loaders.test`. 

287 

288 Parameters 

289 ---------- 

290 variables: dict[str, Any] 

291 Workflow configuration from ROSE_SUITE_VARIABLES. 

292 

293 Returns 

294 ------- 

295 Iterator[RawRecipe] 

296 Configured recipes. 

297 

298 Raises 

299 ------ 

300 AttributeError 

301 When a loader doesn't provide a `load` function. 

302 """ 

303 # Import here to avoid circular import. 

304 import CSET.loaders 

305 

306 config = Config(variables) 

307 for loader in CSET.loaders.__all__: 

308 logger.info("Loading recipes from %s", loader) 

309 module = getattr(CSET.loaders, loader) 

310 yield from module.load(config)