Coverage for src/CSET/operators/write.py: 100%
14 statements
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-01 15:05 +0000
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-01 15:05 +0000
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.
15"""Operators for writing various types of files to disk."""
17import secrets
18from pathlib import Path
19from typing import Union
21import iris
22import iris.cube
24from CSET._common import get_recipe_metadata, slugify
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.
35 This operator expects an iris cube object that will then be saved to disk.
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.
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"))
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"
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