Coverage for /opt/conda/envs/apienv/lib/python3.10/site-packages/daiquiri/core/saving/base.py: 85%

1#!/usr/bin/env python 

2# -*- coding: utf-8 -*- 

3import json 

4import os 

5import logging 

6import re 

7import string 

8from abc import abstractmethod, abstractproperty 

9import time 

10from flask import g 

11from marshmallow import fields 

12 

13from daiquiri.core import CoreBase, CoreResource, marshal 

14from daiquiri.core.utils import loader 

15from daiquiri.core.schema import ErrorSchema, MessageSchema 

16from daiquiri.core.schema.saving import SavingSchema 

17 

18logger = logging.getLogger(__name__) 

19 

20 

21class Saving: 

22 def __init__(self, config, *args, **kwargs): 

23 kwargs["config"] = config 

24 

25 try: 

26 self._saving = loader( 

27 "daiquiri.core.saving", 

28 "SavingHandler", 

29 config["saving_type"], 

30 *args, 

31 **kwargs, 

32 ) 

33 except KeyError: 

34 logger.warning( 

35 "Configuration section for saving missing, saving component not configured" 

36 ) 

37 self._saving = None 

38 

39 def init(self): 

40 return self._saving 

41 

42 

43class SavingResource(CoreResource): 

44 @marshal( 

45 inp={ 

46 "sessionid": fields.Int(), 

47 "sampleid": fields.Int(), 

48 "subsampleid": fields.Int(), 

49 }, 

50 out=[[200, SavingSchema(), "The current saving setup"]], 

51 ) 

52 def get(self, **kwargs): 

53 """Get the current saving configuration""" 

54 kwargs["sessionid"] = g.blsession.get("sessionid") 

55 return {"arguments": self._parent.get_schema_fields(**kwargs)}, 200 

56 

57 @marshal( 

58 inp=SavingSchema, 

59 out=[ 

60 [200, MessageSchema(), "Updated saving arguments"], 

61 [400, ErrorSchema(), "Could not update saving arguments"], 

62 ], 

63 ) 

64 def patch(self, **kwargs): 

65 """Update the saving arguments""" 

66 if False: 

67 return {"message": "saving arguments updated"} 

68 else: 

69 return {"error": "Could not update saving arguments"}, 400 

70 

71 

72class FormatterObject: 

73 """Used in interpolating Saving arguments""" 

74 

75 def __init__(self, resource, srepr, info): 

76 self._resource = resource 

77 self._srepr = srepr 

78 self._info = info 

79 

80 def __str__(self): 

81 return self._srepr 

82 

83 def _make_safe(self, value): 

84 """Make a value safe for the file system 

85 

86 Replaces spaces and special characters with underscore 

87 """ 

88 if value is None: 

89 return "" 

90 

91 if isinstance(value, dict): 

92 return json.dumps(value, indent=2) 

93 else: 

94 return re.sub(r"[^\w\.-]", "_", value) 

95 

96 def __getattr__(self, attr): 

97 try: 

98 return self._make_safe(self._info[attr]) 

99 except KeyError: 

100 raise AttributeError( 

101 f"{self._resource}.{attr} (available: {list(self._info.keys())}" 

102 ) from None 

103 

104 

105class SavingFormatter(string.Formatter): 

106 """Used for interpolating Saving arguments""" 

107 

108 def __init__(self, handler, actor_args, partial_evaulation=False): 

109 """ 

110 :param SavingHandler handler: 

111 """ 

112 super().__init__() 

113 self._handler = handler 

114 self._actor_args = actor_args 

115 self._partial_evaluation = partial_evaulation 

116 

117 def get_value(self, key, args, kwds): 

118 """Order of finding the key: 

119 

120 - format named arguments 

121 - handler config 

122 - actor arguments 

123 """ 

124 if isinstance(key, str): 

125 if key == "epoch": 

126 return str(int(time.time())) 

127 try: 

128 return kwds[key] 

129 except KeyError: 

130 pass 

131 try: 

132 return self._handler._config[key] 

133 except KeyError: 

134 pass 

135 try: 

136 return self._get_from_actor_args(key) 

137 except KeyError: 

138 pass 

139 

140 #  TODO: Which edge cases does this not work for? 

141 if self._partial_evaluation: 

142 return f"{{{key}}}" 

143 

144 return super().get_value(key, args, kwds) 

145 

146 @property 

147 def _metadata(self): 

148 return self._handler._metadata 

149 

150 def _get_from_actor_args(self, key): 

151 """Get key from actor arguments and allow special metadata keys 

152 to be compounded like `{sessionid.name}` 

153 

154 :param str key: 

155 :returns FormatterObject: 

156 """ 

157 # TODO: too specific 

158 if key == "componentid": 

159 value = self._actor_args["sampleid"] 

160 else: 

161 value = self._actor_args[key] 

162 

163 info = None 

164 if key == "sessionid": 

165 info = self._metadata.get_sessions(sessionid=value, no_context=True) 

166 elif key == "sampleid": 

167 info = self._metadata.get_samples(sampleid=value, no_context=True) 

168 elif key == "componentid": 

169 info = self._metadata.get_components(sampleid=value, no_context=True) 

170 if info["rows"]: 

171 info = info["rows"][0] 

172 elif key == "subsampleid": 

173 info = self._metadata.get_subsamples(subsampleid=value, no_context=True) 

174 elif key == "datacollectionid": 

175 info = self._metadata.get_datacollections( 

176 datacollectionid=value, no_context=True 

177 ) 

178 if info is None: 

179 info = {} 

180 return FormatterObject( 

181 key, 

182 json.dumps(value, indent=2) if isinstance(value, dict) else str(value), 

183 info, 

184 ) 

185 

186 

187class SavingHandler(CoreBase): 

188 _require_session = True 

189 _require_blsession = True 

190 _base_url = "saving" 

191 

192 _callbacks = [] 

193 _formatter = string.Formatter() 

194 

195 def setup(self): 

196 self.saving_arguments = self._config.get("saving_arguments", {}) 

197 self.register_route(SavingResource, "") 

198 

199 def eval_saving_arguments( 

200 self, 

201 eval_args, 

202 extra_saving_args=None, 

203 raise_on_missing=True, 

204 warn_on_missing=False, 

205 partial_evaulation=False, 

206 saving_arguments=None, 

207 ): 

208 logger.debug(f"saving arguments {eval_args}") 

209 formatter = SavingFormatter(self, eval_args, partial_evaulation) 

210 

211 if saving_arguments: 

212 all_saving_args = {**saving_arguments} 

213 else: 

214 all_saving_args = {**self.saving_arguments} 

215 if extra_saving_args: 

216 all_saving_args.update(extra_saving_args) 

217 

218 kwargs = {} 

219 for name, template in all_saving_args.items(): 

220 try: 

221 kwargs[name] = formatter.format(template) 

222 except (KeyError, AttributeError): 

223 if raise_on_missing: 

224 raise 

225 #  With warnings enabled the unreplaced key will not be available 

226 elif warn_on_missing: 

227 logger.warning( 

228 f"Could not evaluate template variable {name}: {template}" 

229 ) 

230 else: 

231 kwargs[name] = template 

232 

233 return kwargs 

234 

235 def get_schema_fields(self, **eval_args): 

236 """Return the saving arguments""" 

237 saving_arguments = self.eval_saving_arguments( 

238 eval_args, raise_on_missing=False, partial_evaulation=True 

239 ) 

240 saving_arguments["proposal_root_path"] = self.proposal_root_path 

241 return saving_arguments 

242 

243 def set_filename(self, extra_saving_args=None, set_metadata=True, **eval_args): 

244 """Define the file name""" 

245 self._set_filename( 

246 **self.eval_saving_arguments(eval_args, extra_saving_args=extra_saving_args) 

247 ) 

248 if set_metadata and self._config.get("saving_metadata"): 

249 self.set_metadata(self._config["saving_metadata"], **eval_args) 

250 return self.filename 

251 

252 @abstractmethod 

253 def _set_filename(self, **eval_args): 

254 """Child class implementation of applying the saving arguments""" 

255 pass 

256 

257 @abstractproperty 

258 def filename(self): 

259 """Child class implementation of returning the current filename""" 

260 pass 

261 

262 def set_metadata(self, saving_arguments, warn_on_missing=True, **eval_args): 

263 """Evaluate metadata required for metadata catalogue""" 

264 evaluated = self.eval_saving_arguments( 

265 eval_args, 

266 saving_arguments=saving_arguments, 

267 raise_on_missing=not warn_on_missing, 

268 warn_on_missing=warn_on_missing, 

269 ) 

270 self._set_metadata(**evaluated) 

271 return evaluated 

272 

273 def _set_metadata(self, **metadata): 

274 """Apply the evaluated metadata to whatever service is required 

275 Should be overridden in the child class 

276 """ 

277 pass 

278 

279 @abstractproperty 

280 def create_root_path(self): 

281 pass 

282 

283 @abstractmethod 

284 def create_path(self, path): 

285 pass 

286 

287 @property 

288 def dirname(self): 

289 return os.path.dirname(self.filename) 

290 

291 @property 

292 def basename(self): 

293 return os.path.basename(self.filename) 

294 

295 @property 

296 def proposal_root_path(self): 

297 pass