Coverage for src/CSET/operators/write.py: 100%

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"""Operators for writing various types of files to disk.""" 

16 

17import secrets 

18from pathlib import Path 

19 

20import iris 

21import iris.cube 

22 

23from CSET._common import get_recipe_metadata, slugify 

24 

25 

26def write_cube_to_nc( 

27 cube: iris.cube.Cube | iris.cube.CubeList, 

28 filename: str = None, 

29 overwrite: bool = False, 

30 **kwargs, 

31) -> str: 

32 """Write a cube to a NetCDF file. 

33 

34 This operator expects an iris cube object that will then be saved to disk. 

35 

36 Arguments 

37 --------- 

38 cube: iris.cube.Cube | iris.cube.CubeList 

39 Data to save. 

40 filename: str, optional 

41 Path to save the cubes too. Defaults to the recipe title + .nc 

42 overwrite: bool, optional 

43 Whether to overwrite an existing file. If False the filename will have a 

44 unique suffix added. Defaults to False. 

45 

46 Returns 

47 ------- 

48 Cube | CubeList 

49 The inputted cube(list) (so further operations can be applied) 

50 """ 

51 # Allow writing to be disabled without changing the recipe. This improves 

52 # runtime and avoids using excessive disk space for large runs. 

53 if get_recipe_metadata().get("skip_write"): 

54 return cube 

55 

56 if filename is None: 

57 filename = slugify(get_recipe_metadata().get("title", "Untitled")) 

58 

59 # Append a unique suffix if not overwriting. We use randomness rather than a 

60 # sequence number to avoid race conditions with multiple job runners. 

61 if not overwrite: 

62 filename = f"{filename}_{secrets.token_urlsafe(16)}.nc" 

63 

64 # Ensure that output filename is a Path with a .nc suffix 

65 filename = Path(filename).with_suffix(".nc") 

66 # Save the file as nc compliant (iris should handle this) 

67 iris.save(cube, filename, zlib=True) 

68 return cube