1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ("FormatterTest", "DoNothingFormatter", "LenientYamlFormatter", "MetricsExampleFormatter", 

25 "MultipleExtensionsFormatter", "SingleExtensionFormatter") 

26 

27from typing import ( 

28 TYPE_CHECKING, 

29 Any, 

30 Mapping, 

31 Optional, 

32) 

33 

34import json 

35import yaml 

36 

37from ..core import Formatter 

38from ..formatters.yaml import YamlFormatter 

39 

40if TYPE_CHECKING: 40 ↛ 41line 40 didn't jump to line 41, because the condition on line 40 was never true

41 from ..core import Location 

42 

43 

44class DoNothingFormatter(Formatter): 

45 """A test formatter that does not need to format anything and has 

46 parameters.""" 

47 

48 def read(self, component: Optional[str] = None) -> Any: 

49 raise NotImplementedError("Type does not support reading") 

50 

51 def write(self, inMemoryDataset: Any) -> str: 

52 raise NotImplementedError("Type does not support writing") 

53 

54 

55class FormatterTest(Formatter): 

56 """A test formatter that does not need to format anything.""" 

57 

58 supportedWriteParameters = frozenset({"min", "max", "median", "comment", "extra", "recipe"}) 

59 

60 def read(self, component: Optional[str] = None) -> Any: 

61 raise NotImplementedError("Type does not support reading") 

62 

63 def write(self, inMemoryDataset: Any) -> str: 

64 raise NotImplementedError("Type does not support writing") 

65 

66 @staticmethod 

67 def validateWriteRecipes(recipes: Optional[Mapping[str, Any]]) -> Optional[Mapping[str, Any]]: 

68 if not recipes: 

69 return recipes 

70 for recipeName in recipes: 

71 if "mode" not in recipes[recipeName]: 

72 raise RuntimeError("'mode' is a required write recipe parameter") 

73 return recipes 

74 

75 

76class SingleExtensionFormatter(DoNothingFormatter): 

77 """A do nothing formatter that has a single extension registered.""" 

78 extension = ".fits" 

79 

80 

81class MultipleExtensionsFormatter(SingleExtensionFormatter): 

82 """A formatter that has multiple extensions registered.""" 

83 supportedExtensions = frozenset({".fits.gz", ".fits.fz", ".fit"}) 

84 

85 

86class LenientYamlFormatter(YamlFormatter): 

87 """A test formatter that allows any file extension but always reads and 

88 writes YAML.""" 

89 extension = ".yaml" 

90 

91 @classmethod 

92 def validateExtension(cls, location: Location) -> None: 

93 return 

94 

95 

96class MetricsExampleFormatter(Formatter): 

97 """A specialist test formatter for metrics that supports components 

98 directly without assembler delegate.""" 

99 

100 supportedExtensions = frozenset({".yaml", ".json"}) 

101 

102 @property 

103 def extension(self) -> str: 

104 """Always write yaml by default.""" 

105 return ".yaml" 

106 

107 def read(self, component=None): 

108 """Read data from a file. 

109 

110 Parameters 

111 ---------- 

112 component : `str`, optional 

113 Component to read from the file. Only used if the `StorageClass` 

114 for reading differed from the `StorageClass` used to write the 

115 file. 

116 

117 Returns 

118 ------- 

119 inMemoryDataset : `object` 

120 The requested data as a Python object. The type of object 

121 is controlled by the specific formatter. 

122 

123 Raises 

124 ------ 

125 ValueError 

126 Component requested but this file does not seem to be a concrete 

127 composite. 

128 KeyError 

129 Raised when parameters passed with fileDescriptor are not 

130 supported. 

131 """ 

132 

133 # This formatter can not read a subset from disk because it 

134 # uses yaml or json. 

135 path = self.fileDescriptor.location.path 

136 

137 with open(path, "r") as fd: 

138 if path.endswith(".yaml"): 

139 data = yaml.load(fd, Loader=yaml.SafeLoader) 

140 elif path.endswith(".json"): 

141 data = json.load(fd) 

142 else: 

143 raise RuntimeError(f"Unsupported file extension found in path '{path}'") 

144 

145 # We can slice up front if required 

146 parameters = self.fileDescriptor.parameters 

147 if "data" in data and parameters and "slice" in parameters: 

148 data["data"] = data["data"][parameters["slice"]] 

149 

150 pytype = self.fileDescriptor.storageClass.pytype 

151 inMemoryDataset = pytype(**data) 

152 

153 if not component: 

154 return inMemoryDataset 

155 

156 if component == "summary": 

157 return inMemoryDataset.summary 

158 elif component == "output": 

159 return inMemoryDataset.output 

160 elif component == "data": 

161 return inMemoryDataset.data 

162 elif component == "counter": 

163 return len(inMemoryDataset.data) 

164 raise ValueError(f"Unsupported component: {component}") 

165 

166 def write(self, inMemoryDataset: Any) -> str: 

167 """Write a Dataset. 

168 

169 Parameters 

170 ---------- 

171 inMemoryDataset : `object` 

172 The Dataset to store. 

173 

174 Returns 

175 ------- 

176 path : `str` 

177 The path to where the Dataset was stored within the datastore. 

178 """ 

179 fileDescriptor = self.fileDescriptor 

180 

181 # Update the location with the formatter-preferred file extension 

182 fileDescriptor.location.updateExtension(self.extension) 

183 

184 with open(fileDescriptor.location.path, "w") as fd: 

185 yaml.dump(inMemoryDataset._asdict(), fd) 

186 return fileDescriptor.location.pathInStore 

187 

188 

189class MetricsExampleDataFormatter(Formatter): 

190 """A specialist test formatter for the data component of a MetricsExample. 

191 

192 This is needed if the MetricsExample is dissassembled and we want to 

193 support the derived component. 

194 """ 

195 

196 unsupportedParameters = None 

197 """Let the assembler delegate handle slice""" 

198 

199 extension = ".yaml" 

200 """Always write YAML""" 

201 

202 def read(self, component=None): 

203 """Read data from a file. 

204 

205 Parameters 

206 ---------- 

207 component : `str`, optional 

208 Component to read from the file. Only used if the `StorageClass` 

209 for reading differed from the `StorageClass` used to write the 

210 file. 

211 

212 Returns 

213 ------- 

214 inMemoryDataset : `object` 

215 The requested data as a Python object. The type of object 

216 is controlled by the specific formatter. 

217 

218 Raises 

219 ------ 

220 ValueError 

221 Component requested but this file does not seem to be a concrete 

222 composite. 

223 KeyError 

224 Raised when parameters passed with fileDescriptor are not 

225 supported. 

226 """ 

227 

228 # This formatter can not read a subset from disk because it 

229 # uses yaml. 

230 path = self.fileDescriptor.location.path 

231 with open(path, "r") as fd: 

232 data = yaml.load(fd, Loader=yaml.SafeLoader) 

233 

234 # We can slice up front if required 

235 parameters = self.fileDescriptor.parameters 

236 if parameters and "slice" in parameters: 

237 data = data[parameters["slice"]] 

238 

239 # This should be a native list 

240 inMemoryDataset = data 

241 

242 if not component: 

243 return inMemoryDataset 

244 

245 if component == "counter": 

246 return len(inMemoryDataset) 

247 raise ValueError(f"Unsupported component: {component}") 

248 

249 def write(self, inMemoryDataset: Any) -> str: 

250 """Write a Dataset. 

251 

252 Parameters 

253 ---------- 

254 inMemoryDataset : `object` 

255 The Dataset to store. 

256 

257 Returns 

258 ------- 

259 path : `str` 

260 The path to where the Dataset was stored within the datastore. 

261 """ 

262 fileDescriptor = self.fileDescriptor 

263 

264 # Update the location with the formatter-preferred file extension 

265 fileDescriptor.location.updateExtension(self.extension) 

266 

267 with open(fileDescriptor.location.path, "w") as fd: 

268 yaml.dump(inMemoryDataset, fd) 

269 return fileDescriptor.location.pathInStore