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

1# Copyright 2022 Met Office and 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 

19from typing import Union 

20 

21import iris 

22import iris.cube 

23 

24from CSET._common import get_recipe_metadata, slugify 

25 

26 

27def write_cube_to_nc( 

28 cube: Union[iris.cube.Cube, iris.cube.CubeList], 

29 filename: str = None, 

30 overwrite: bool = False, 

31 **kwargs, 

32) -> str: 

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

34 

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

36 

37 Arguments 

38 --------- 

39 cube: iris.cube.Cube | iris.cube.CubeList 

40 Data to save. 

41 filename: str, optional 

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

43 overwrite: bool, optional 

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

45 unique suffix added. Defaults to False. 

46 

47 Returns 

48 ------- 

49 Cube | CubeList 

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

51 """ 

52 if filename is None: 

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

54 

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

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

57 if not overwrite: 

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

59 

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

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

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

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

64 return cube