Coverage for /opt/conda/envs/apienv/lib/python3.10/site-packages/daiquiri/core/components/parameteriser.py: 90%

1#!/usr/bin/env python 

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

3import glob 

4import os 

5from flask import g 

6from marshmallow import fields 

7import ruamel 

8 

9from daiquiri.core import marshal 

10from daiquiri.core.components import ( 

11 Component, 

12 ComponentActor, 

13 ComponentActorSchema, 

14 ComponentResource, 

15 actor, 

16) 

17from daiquiri.core.schema import ErrorSchema 

18from daiquiri.core.schema.validators import ValidatedRegexp 

19from daiquiri.core.schema.metadata import paginated 

20from daiquiri.core.schema.parameteriser import ( 

21 ParameteriserConfigSchema, 

22 AvailableParametersSchema, 

23) 

24 

25try: 

26 from bliss.scanning.scan_saving import ESRFScanSaving 

27except ImportError: 

28 ESRFScanSaving = None 

29 

30 

31import logging 

32 

33 

34logger = logging.getLogger(__name__) 

35 

36 

37class AvailableParametersResource(ComponentResource): 

38 @marshal( 

39 inp={ 

40 "type": fields.Str( 

41 required=True, 

42 metadata={"description": "The type of parameters to list"}, 

43 ) 

44 }, 

45 out=[ 

46 [200, paginated(AvailableParametersSchema), "List of parameters"], 

47 [400, ErrorSchema(), "Could not get list of parameters"], 

48 ], 

49 ) 

50 def get(self, type: str): 

51 """Get the list of parameter sets for the selected parameter type""" 

52 try: 

53 params = self._parent.get_available_parameters(type) 

54 return {"total": len(params), "rows": params}, 200 

55 except Exception: 

56 message = f"Could not get parameters for type: {type}" 

57 logger.exception(message) 

58 return {"error": message}, 400 

59 

60 

61class Parameteriser(Component): 

62 _base_url = "parameteriser" 

63 _config_schema = ParameteriserConfigSchema() 

64 _config_export = ["parametertypes"] 

65 

66 def setup(self): 

67 self._parameters_actors = [] 

68 self.register_route(AvailableParametersResource, "") 

69 self._generate_parameter_actors() 

70 

71 def reload(self): 

72 self._generate_parameter_actors() 

73 

74 def _generate_parameter_actors(self): 

75 """Dynamically generate parameter actor resources""" 

76 

77 def post(self, **kwargs): 

78 pass 

79 

80 def preprocess(self, **kwargs): 

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

82 

83 return kwargs 

84 

85 if not self._config.get("actors"): 

86 self._config["actors"] = {} 

87 

88 for _, actors in self._config["parametertypes"].items(): 

89 for actor_desc in actors: 

90 actor_name = actor_desc["actor"] 

91 

92 self._config["actors"][actor_name] = actor_name 

93 

94 if actor_name in self._actors: 

95 continue 

96 

97 self._actors.append(actor_name) 

98 self._parameters_actors.append(actor_name) 

99 

100 fn = actor(actor_name, spawn=True, preprocess=True)(post) 

101 

102 act_res = type( 

103 actor_name, 

104 (ComponentResource,), 

105 {"post": fn, "preprocess": preprocess}, 

106 ) 

107 self.register_actor_route(act_res, f"/{actor_name}") 

108 

109 def set_proposal(self, **data): 

110 args = self._saving.eval_saving_arguments( 

111 data, partial_evaulation=True, raise_on_missing=False 

112 ) 

113 

114 arg_map = { 

115 "proposal": "proposal_name", 

116 "sample": "collection_name", 

117 "dataset": "dataset_name", 

118 } 

119 

120 for arg, value in args.items(): 

121 # dont set unevaluated keys (we really only want to set the base path) 

122 if "{" not in value: 

123 key = arg 

124 

125 # For ESRF data policy saving we have do some arg mapping as per: 

126 # daiquiri/core/saving/bliss_esrf.py 

127 if arg in arg_map: 

128 key = arg_map[arg] 

129 

130 setattr(self._saving.scan_saving, key, value) 

131 

132 @property 

133 def root_directory(self): 

134 """Return the root proposal directory""" 

135 # TODO: this is a bit bliss specific 

136 if isinstance(self._saving.scan_saving, ESRFScanSaving): 

137 return os.path.join( 

138 self._saving.scan_saving.base_path, 

139 self._saving.scan_saving.proposal_name, 

140 self._saving.scan_saving.beamline, 

141 ) 

142 else: 

143 return self._saving.scan_saving.base_path 

144 

145 def actor_started(self, actid, actor): 

146 """Callback when an actor starts 

147 

148 Just set the correct proposal 

149 """ 

150 if actor.name in self._parameters_actors: 

151 self.set_proposal(**actor.all_data) 

152 actor.update( 

153 root_directory=self.root_directory, 

154 parameter_root=self._config.get("root", "parameteriser"), 

155 ) 

156 

157 logger.info(f"Actor '{actor.name}' with id '{actid}' started") 

158 

159 def get_available_parameters(self, type): 

160 self.set_proposal(**{"sessionid": g.blsession.get("sessionid")}) 

161 parameter_path = os.path.join( 

162 self.root_directory, self._config.get("root", "parameteriser"), type 

163 ) 

164 files = glob.glob(f"{parameter_path}/*.yml") 

165 

166 parameters = [] 

167 for file in files: 

168 with open(file) as stream: 

169 loader = ruamel.yaml.Loader(stream) 

170 data = loader.get_single_data() 

171 print("param", file, data) 

172 parameters.append( 

173 { 

174 "parameter_type": data.pop("parameter_type"), 

175 "instance_type": data.pop("instance_type"), 

176 "name": data["name"], 

177 "file_name": os.path.basename(file), 

178 "directory": os.path.dirname(file), 

179 "parameters": data, 

180 } 

181 ) 

182 

183 return parameters 

184 

185 

186class ParameteriserActorSchema(ComponentActorSchema): 

187 name = ValidatedRegexp("word-dash", required=True, metadata={"title": "Name"}) 

188 overwrite = fields.Bool() 

189 

190 def __init__(self, *args, **kwargs): 

191 super().__init__(*args, **kwargs) 

192 

193 self.Meta.cache = False 

194 

195 if not hasattr(self.Meta, "uischema"): 

196 self.Meta.uischema = {} 

197 

198 self.Meta.uischema["overwrite"] = { 

199 "classNames": "hidden-row", 

200 "ui:widget": "hidden", 

201 } 

202 

203 

204class ParameteriserActor(ComponentActor): 

205 def method(self, name, overwrite=False, **kwargs): 

206 parameter_path = os.path.join( 

207 self["root_directory"], 

208 self["parameter_root"], 

209 self.parameter_type, 

210 f"{self.name}_{name}.yml", 

211 ) 

212 

213 if os.path.exists(parameter_path) and not overwrite: 

214 raise RuntimeError( 

215 f"Parameter file `{parameter_path}` already exists and not overwriting" 

216 ) 

217 

218 base_dir = os.path.dirname(parameter_path) 

219 if not os.path.exists(base_dir): 

220 os.makedirs(base_dir) 

221 

222 all_kwargs = { 

223 "parameter_type": self.parameter_type, 

224 "instance_type": self.name, 

225 "name": name, 

226 **kwargs, 

227 } 

228 

229 yaml = ruamel.yaml.YAML() 

230 yaml.allow_duplicate_keys = True 

231 yaml.default_flow_style = False 

232 yaml.indent(mapping=2, sequence=4, offset=2) 

233 

234 with open(parameter_path, "wt") as parameter_file: 

235 yaml.dump(all_kwargs, stream=parameter_file)