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

1# Copyright 2022-2023 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"""CSET: Convective and turbulence Scale Evaluation Tool.""" 

16 

17import argparse 

18import logging 

19import os 

20import shlex 

21import sys 

22from importlib.metadata import version 

23from pathlib import Path 

24 

25from CSET._common import ArgumentError 

26 

27 

28def main(): 

29 """CLI entrypoint.""" 

30 parser = argparse.ArgumentParser( 

31 prog="cset", description="Convective Scale Evaluation Tool" 

32 ) 

33 parser.add_argument( 

34 "-v", 

35 "--verbose", 

36 action="count", 

37 default=0, 

38 help="increase output verbosity, may be specified multiple times", 

39 ) 

40 parser.add_argument( 

41 "--version", action="version", version=f"CSET v{version('CSET')}" 

42 ) 

43 

44 # https://docs.python.org/3/library/argparse.html#sub-commands 

45 subparsers = parser.add_subparsers(title="subcommands", dest="subparser") 

46 

47 # Run operator chain 

48 parser_bake = subparsers.add_parser("bake", help="run steps from a recipe file") 

49 parser_bake.add_argument( 

50 "-i", 

51 "--input-dir", 

52 type=Path, 

53 help="directory containing input data", 

54 ) 

55 parser_bake.add_argument( 

56 "-o", 

57 "--output-dir", 

58 type=Path, 

59 required=True, 

60 help="directory to write output into", 

61 ) 

62 parser_bake.add_argument( 

63 "-r", 

64 "--recipe", 

65 type=Path, 

66 required=True, 

67 help="recipe file to read", 

68 ) 

69 bake_step_control = parser_bake.add_mutually_exclusive_group() 

70 bake_step_control.add_argument( 

71 "--parallel-only", action="store_true", help="only run parallel steps" 

72 ) 

73 bake_step_control.add_argument( 

74 "--collate-only", action="store_true", help="only run collation steps" 

75 ) 

76 parser_bake.add_argument( 

77 "-s", "--style-file", type=Path, help="colour bar definition to use" 

78 ) 

79 parser_bake.set_defaults(func=_bake_command) 

80 

81 parser_graph = subparsers.add_parser("graph", help="visualise a recipe file") 

82 parser_graph.add_argument( 

83 "-d", 

84 "--details", 

85 action="store_true", 

86 help="include operator arguments in output", 

87 ) 

88 parser_graph.add_argument( 

89 "-o", 

90 "--output-path", 

91 type=Path, 

92 nargs="?", 

93 help="persistent file to save the graph. Otherwise the file is opened", 

94 default=None, 

95 ) 

96 parser_graph.add_argument( 

97 "-r", 

98 "--recipe", 

99 type=Path, 

100 required=True, 

101 help="recipe file to read", 

102 ) 

103 parser_graph.set_defaults(func=_graph_command) 

104 

105 parser_cookbook = subparsers.add_parser( 

106 "cookbook", help="unpack included recipes to a folder" 

107 ) 

108 parser_cookbook.add_argument( 

109 "-d", 

110 "--details", 

111 action="store_true", 

112 help="list available recipes. Supplied recipes are detailed", 

113 ) 

114 parser_cookbook.add_argument( 

115 "-o", 

116 "--output-dir", 

117 type=Path, 

118 help="directory to save recipes. If omitted uses $PWD", 

119 default=Path.cwd(), 

120 ) 

121 parser_cookbook.add_argument( 

122 "recipe", 

123 type=str, 

124 nargs="?", 

125 help="recipe to output or detail", 

126 default="", 

127 ) 

128 parser_cookbook.set_defaults(func=_cookbook_command) 

129 

130 parser_recipe_id = subparsers.add_parser("recipe-id", help="get the ID of a recipe") 

131 parser_recipe_id.add_argument( 

132 "-r", 

133 "--recipe", 

134 type=Path, 

135 required=True, 

136 help="recipe file to read", 

137 ) 

138 parser_recipe_id.set_defaults(func=_recipe_id_command) 

139 

140 cli_args = sys.argv[1:] + shlex.split(os.getenv("CSET_ADDOPTS", "")) 

141 args, unparsed_args = parser.parse_known_args(cli_args) 

142 

143 # Setup logging. 

144 logging.captureWarnings(True) 

145 loglevel = calculate_loglevel(args) 

146 logger = logging.getLogger() 

147 logger.setLevel(min(loglevel, logging.INFO)) 

148 stderr_log = logging.StreamHandler() 

149 stderr_log.addFilter(lambda record: record.levelno >= loglevel) 

150 stderr_log.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s")) 

151 logger.addHandler(stderr_log) 

152 

153 if args.subparser is None: 

154 print("Please choose a command.", file=sys.stderr) 

155 parser.print_usage() 

156 sys.exit(127) 

157 

158 try: 

159 # Execute the specified subcommand. 

160 args.func(args, unparsed_args) 

161 except ArgumentError as err: 

162 logging.error(err) 

163 parser.print_usage() 

164 sys.exit(3) 

165 

166 

167def calculate_loglevel(args) -> int: 

168 """Calculate the logging level to apply. 

169 

170 Level is based on verbose argument and the LOGLEVEL environment variable. 

171 """ 

172 try: 

173 name_to_level = logging.getLevelNamesMapping() 

174 except AttributeError: 

175 # logging.getLevelNamesMapping() is python 3.11 or newer. Using 

176 # implementation detail for older versions. 

177 name_to_level = logging._nameToLevel 

178 # Level from CLI flags. 

179 if args.verbose >= 2: 

180 loglevel = logging.DEBUG 

181 elif args.verbose == 1: 

182 loglevel = logging.INFO 

183 else: 

184 loglevel = logging.WARNING 

185 return min( 

186 loglevel, 

187 # Level from environment variable. 

188 name_to_level.get(os.getenv("LOGLEVEL"), logging.ERROR), 

189 ) 

190 

191 

192def _bake_command(args, unparsed_args): 

193 from CSET._common import parse_variable_options 

194 from CSET.operators import execute_recipe_collate, execute_recipe_parallel 

195 

196 recipe_variables = parse_variable_options(unparsed_args) 

197 if not args.collate_only: 

198 # Input dir is needed for parallel steps, but not collate steps. 

199 if not args.input_dir: 

200 raise ArgumentError("the following arguments are required: -i/--input-dir") 

201 execute_recipe_parallel( 

202 args.recipe, 

203 args.input_dir, 

204 args.output_dir, 

205 recipe_variables, 

206 args.style_file, 

207 ) 

208 if not args.parallel_only: 

209 execute_recipe_collate( 

210 args.recipe, args.output_dir, recipe_variables, args.style_file 

211 ) 

212 

213 

214def _graph_command(args, unparsed_args): 

215 from CSET.graph import save_graph 

216 

217 save_graph( 

218 args.recipe, 

219 args.output_path, 

220 auto_open=not args.output_path, 

221 detailed=args.details, 

222 ) 

223 

224 

225def _cookbook_command(args, unparsed_args): 

226 from CSET.recipes import detail_recipe, list_available_recipes, unpack_recipe 

227 

228 if args.recipe: 

229 if args.details: 

230 detail_recipe(args.recipe) 

231 else: 

232 try: 

233 unpack_recipe(args.output_dir, args.recipe) 

234 except FileNotFoundError: 

235 logging.error("Recipe %s does not exist.", args.recipe) 

236 sys.exit(1) 

237 else: 

238 list_available_recipes() 

239 

240 

241def _recipe_id_command(args, unparsed_args): 

242 from uuid import uuid4 

243 

244 from CSET._common import parse_recipe, parse_variable_options, slugify 

245 

246 recipe_variables = parse_variable_options(unparsed_args) 

247 recipe = parse_recipe(args.recipe, recipe_variables) 

248 try: 

249 recipe_id = slugify(recipe["title"]) 

250 except KeyError: 

251 logging.warning("Recipe has no title; Falling back to random recipe_id.") 

252 recipe_id = str(uuid4()) 

253 print(recipe_id)