Coverage for src/susi/base/loader.py: 83%
109 statements
« prev ^ index » next coverage.py v7.5.0, created at 2025-08-22 09:20 +0000
1"""
2Loader module provides the Loder class for scripts
4@author hoelken@mps.mpg.de
5"""
7import os
8from datetime import datetime
9from argparse import Namespace, ArgumentParser
11from .exceptions import MissConfigurationException
12from .config import Config
13from .logging import Logging
14from ..utils.git import Git
15from ..utils.yaml import write, read_yaml
16from .api import Api
18log = Logging.get_logger()
21class Loader:
22 """Helper class for scripts to DRY out common setup tasks."""
24 def __init__(self, args: Namespace, script: str, yaml_config: dict = None):
25 self.args = args
26 self.yaml = yaml_config if yaml_config is not None else read_yaml(args.conf)
27 self.script = script
28 self.log_dir = None
29 self.config = None
30 self.start = None
31 self.stop = None
32 self.dim_keys = None
34 def setup(self, git: bool = True):
35 """
36 Setup the configuration from the command line arguments
37 """
38 if self.args.verbose:
39 Logging.set_log_level(0)
40 for c in ['cam1', 'cam2', 'cam3']:
41 if c in self.yaml:
42 if 'data' in self.yaml[c]:
43 if 'log_dir' in self.yaml[c]['data']:
44 self.log_dir = os.path.join(
45 self.yaml[c]['data']['log_dir'],
46 'susi_reduction_' + datetime.now().strftime("%Y%m%d"),
47 datetime.now().strftime("%H%M%S"),
48 )
50 if self.args.log_dir:
51 self.log_dir = os.path.join(
52 self.args.log_dir, 'susi_datareduction_' + datetime.now().strftime("%y%m%dT%H%M%S")
53 )
55 if self.args.log_fulldir:
56 self.log_dir = self.args.log_fulldir
58 if self.log_dir is not None:
59 Logging.init_file(os.path.join(self.log_dir, f'{self.script}.log'))
61 version = Git.version() if git else None
62 Logging.welcome(self.args, 'SUSI Datareduction', version)
63 if 'config' in self.yaml:
64 if 'cam' in self.yaml['config'] and 'name' in self.yaml['config']['cam']:
65 c = self.yaml['config']['cam']['name']
66 else:
67 c = 'cam1'
68 self.config = Config.from_dict(self.yaml['config'], c)
69 for c in ['cam1', 'cam2', 'cam3']:
70 if c in self.yaml:
71 self.config = Config.from_dict(self.yaml[c], c)
72 if self.log_dir:
73 log.info('Logs and plots will be written to %s', os.path.abspath(self.log_dir))
74 self.config.data.log_dir = self.log_dir
75 log.info('Nice level set to: %s', os.nice(self.config.base.niceness))
76 log.info("======================================")
77 if self.args.id is None:
78 self.start = datetime.fromisoformat(self.args.start if self.args.start is not None else self.yaml['start'])
79 self.stop = datetime.fromisoformat(self.args.stop if self.args.stop is not None else self.yaml['stop'])
80 else:
81 self.start, self.stop = Api(self.config).get_times(self.args.id)
82 self.config.base.obsid = self.args.id
83 self.yaml['start'] = self.start.isoformat()
84 self.config.start = self.start
85 self.yaml['stop'] = self.stop.isoformat()
86 self.config.stop = self.stop
87 if self.log_dir:
88 write(os.path.join(self.log_dir, 'config.yaml'), self.yaml)
89 log.info('DATA START: %s', self.yaml['start'])
90 log.info('DATA STOP: %s', self.yaml['stop'])
91 log.info("======================================")
92 return self
94 def check(self):
95 """Basic configuration check to die early on simple mistakes"""
96 log.debug('Checking for existence of calibration files.')
97 calib_data = self.config.calib_data
98 if calib_data.dark is None and not self.config.base.no_dark_correction:
99 log.critical('No dark file supplied. Set "no_dark_correction: True" to skip dark correction')
100 raise MissConfigurationException('No dark file supplied')
101 elif calib_data.dark is not None and not os.path.isfile(calib_data.dark):
102 log.critical('Dark file not found: %s', calib_data.dark)
103 raise MissConfigurationException("Can't access dark file")
104 if calib_data.mod_matrix is not None and not os.path.isfile(calib_data.mod_matrix):
105 log.critical('Modulation matrix file not found: %s', calib_data.mod_matrix)
106 raise MissConfigurationException("Can't access modulation matrix file")
107 if calib_data.offset_map is not None and not os.path.isfile(calib_data.offset_map):
108 log.critical('Smile map file not found: %s', calib_data.offset_map)
109 raise MissConfigurationException("Can't access smile map file")
110 if calib_data.slit_flat is not None and not os.path.isfile(calib_data.slit_flat):
111 log.critical('Flat field file not found: %s', calib_data.slit_flat)
112 raise MissConfigurationException("Can't access slit flat field file")
113 if calib_data.sensor_flat is not None and not os.path.isfile(calib_data.sensor_flat):
114 log.critical('Flat field file not found: %s', calib_data.sensor_flat)
115 raise MissConfigurationException("Can't access sensor flat field file")
116 self._test_out_path()
117 return self
119 def _test_out_path(self):
120 """Check if we can create files in the output path"""
121 path = os.path.abspath(self.config.out_path())
122 testfile = os.path.join(path, '.test')
123 try:
124 with open(testfile, 'w'):
125 pass
126 except IOError:
127 log.critical('Desired output path is not write-able: %s', path)
128 raise MissConfigurationException("Can't write to out path")
129 finally:
130 if os.path.exists(testfile):
131 os.remove(testfile)
133 @staticmethod
134 def default_arguments(descr: str) -> ArgumentParser:
135 """Provides an `ArgumentParser` with the common CLI arguments already configured."""
136 aparser = ArgumentParser(description=descr)
137 aparser.add_argument('conf', type=str, help='Path to the config file')
138 aparser.add_argument(
139 "-v", "--verbose", dest="verbose", action="store_true", help="Activate verbose mode (DEBUG log level)"
140 )
141 aparser.add_argument(
142 "--log",
143 dest="log_dir",
144 nargs="?",
145 default=False,
146 metavar="DIR",
147 type=str,
148 help="Path to log dir. Timestamped subdir is created. Overwrites log_dir in config file.",
149 )
151 aparser.add_argument(
152 "--log_fulldir",
153 dest="log_fulldir",
154 nargs="?",
155 default=False,
156 metavar="DIR",
157 type=str,
158 help="Path to log dir. NO subdir created. Overwrites --log arg and log_dir in config file.",
159 )
161 aparser.add_argument(
162 "--start",
163 dest="start",
164 nargs="?",
165 metavar="TIME",
166 type=str,
167 default=None,
168 help="Start time in ISO8601 (overwrites start_time in the YAML).",
169 )
170 aparser.add_argument(
171 "--stop",
172 dest="stop",
173 nargs="?",
174 metavar="TIME",
175 type=str,
176 default=None,
177 help="Stop time in ISO8601 (overwrites stop_time in the YAML).",
178 )
179 aparser.add_argument(
180 "--id",
181 dest="id",
182 nargs="?",
183 metavar="INT",
184 type=str,
185 help="ID from the observation log (overwrites all other start/stop time config)",
186 )
187 return aparser