Coverage for python/lsst/images/serialization/_asdf_utils.py: 68%

1# This file is part of lsst-images. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

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

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

7# for details of code ownership. 

8# 

9# Use of this source code is governed by a 3-clause BSD-style 

10# license that can be found in the LICENSE file. 

11 

12from __future__ import annotations 

13 

14__all__ = ( 

15 "ArrayReferenceModel", 

16 "ArrayReferenceQuantityModel", 

17 "InlineArray", 

18 "InlineArrayModel", 

19 "InlineArrayQuantity", 

20 "InlineArrayQuantityModel", 

21 "Quantity", 

22 "QuantityModel", 

23 "Time", 

24 "TimeModel", 

25 "Unit", 

26) 

27 

28from typing import Annotated, Any, Literal 

29 

30import astropy.time 

31import astropy.units 

32import numpy as np 

33import pydantic 

34import pydantic_core.core_schema as pcs 

35from pydantic.json_schema import GetJsonSchemaHandler, JsonSchemaValue 

36 

37from ._dtypes import NumberType 

38 

39 

40class _UnitSerialization: 

41 """Pydantic hooks for unit serialization. 

42 

43 This class provides implementations for the `Unit` type alias for 

44 `astropy.unit.Unit` that adds Pydantic serialization and validation. 

45 """ 

46 

47 @classmethod 

48 def __get_pydantic_core_schema__( 

49 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

50 ) -> pcs.CoreSchema: 

51 from_str_schema = pcs.chain_schema( 

52 [ 

53 pcs.str_schema(), 

54 pcs.no_info_plain_validator_function(cls.from_str), 

55 ] 

56 ) 

57 return pcs.json_or_python_schema( 

58 json_schema=from_str_schema, 

59 python_schema=pcs.union_schema([pcs.is_instance_schema(astropy.units.UnitBase), from_str_schema]), 

60 serialization=pcs.plain_serializer_function_ser_schema(cls.to_str), 

61 ) 

62 

63 @classmethod 

64 def from_str(cls, value: str) -> astropy.units.UnitBase: 

65 try: 

66 return astropy.units.Unit(value, format="vounit") 

67 except ValueError: 

68 pass 

69 # Some important units (e.g. "dn") are not supported by vounit, so 

70 # fall back to letting astropy to infer the format. 

71 return astropy.units.Unit(value) 

72 

73 @staticmethod 

74 def to_str(unit: astropy.units.UnitBase) -> str: 

75 try: 

76 return unit.to_string("vounit") 

77 except ValueError: 

78 pass 

79 # Some important units (e.g. "dn") are not supported by vounit, so 

80 # fall back to letting astropy use the default format. 

81 return unit.to_string() 

82 

83 

84type Unit = Annotated[ 

85 astropy.units.UnitBase, 

86 _UnitSerialization, 

87 pydantic.WithJsonSchema( 

88 { 

89 "type": "string", 

90 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

91 "id": "http://stsci.edu/schemas/asdf/unit/unit-1.0.0", 

92 "tag": "!unit/unit-1.0.0", 

93 } 

94 ), 

95] 

96 

97 

98class ArrayReferenceModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

99 """Model for a subset of the ASDF 'ndarray' schema, in the case where the 

100 array data is stored elsewhere. 

101 """ 

102 

103 source: str | int = pydantic.Field(description="Location of the underlying binary data.") 

104 shape: list[int] = pydantic.Field(description="Size of the array in each dimension.") 

105 datatype: NumberType = pydantic.Field(description="Data type of the array.") 

106 byteorder: Literal["big"] = pydantic.Field(default="big", description="Byte order for the binary data.") 

107 

108 def with_units(self, unit: astropy.units.UnitBase) -> ArrayReferenceQuantityModel: 

109 """Add units, transforming this model into a Quantity model.""" 

110 return ArrayReferenceQuantityModel.model_construct(value=self, unit=unit) 

111 

112 model_config = pydantic.ConfigDict( 

113 json_schema_extra={ 

114 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

115 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0", 

116 "tag": "!core/ndarray-1.1.0", 

117 } 

118 ) 

119 

120 

121class InlineArrayModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

122 """Model for a subset of the ASDF 'ndarray' schema, in the case where the 

123 array data is stored inline. 

124 """ 

125 

126 data: list[Any] 

127 datatype: NumberType 

128 

129 @property 

130 def shape(self) -> tuple[int, ...]: 

131 """The shape of the array (`tuple` [`int`, ...]).""" 

132 return self._extract_shape(self.data) 

133 

134 def with_units(self, unit: astropy.unit.UnitBase) -> InlineArrayQuantityModel: 

135 """Add units, transforming this model in to a Quantity model.""" 

136 return InlineArrayQuantityModel.model_construct(value=self, unit=unit) 

137 

138 @classmethod 

139 def _extract_shape(cls, data: list[Any], current: tuple[int, ...] = ()) -> tuple[int, ...]: 

140 if not data: 

141 return current + (0,) 

142 if not isinstance(data[0], list): 

143 return current + (len(data),) 

144 return cls._extract_shape(data[0], current + (len(data),)) 

145 

146 model_config = pydantic.ConfigDict( 

147 json_schema_extra={ 

148 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

149 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0", 

150 "tag": "!core/ndarray-1.1.0", 

151 } 

152 ) 

153 

154 

155class _InlineArraySerialization: 

156 """Pydantic hooks for array serialization. 

157 

158 This class provides implementations for the `Array` type alias for 

159 `numpy.ndarray` that adds Pydantic serialization and validation. 

160 """ 

161 

162 @classmethod 

163 def __get_pydantic_core_schema__( 

164 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

165 ) -> pcs.CoreSchema: 

166 from_model_schema = pcs.chain_schema( 

167 [ 

168 handler(InlineArrayModel), 

169 pcs.no_info_plain_validator_function(cls.from_model), 

170 ] 

171 ) 

172 return pcs.json_or_python_schema( 

173 json_schema=from_model_schema, 

174 python_schema=pcs.union_schema([pcs.is_instance_schema(np.ndarray), from_model_schema]), 

175 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model), 

176 ) 

177 

178 @classmethod 

179 def __get_pydantic_json_schema__( 

180 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler 

181 ) -> JsonSchemaValue: 

182 return handler(InlineArrayModel.__pydantic_core_schema__) 

183 

184 @classmethod 

185 def from_model(cls, model: InlineArrayModel) -> np.ndarray: 

186 return np.array(model.data, dtype=model.datatype.to_numpy()) 

187 

188 @classmethod 

189 def to_model(cls, array: np.ndarray) -> InlineArrayModel: 

190 datatype = NumberType.from_numpy(array.dtype) 

191 return InlineArrayModel(data=array.tolist(), datatype=datatype) 

192 

193 

194type InlineArray = Annotated[np.ndarray, _InlineArraySerialization] 

195 

196 

197class QuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

198 """Model for a subset of the ASDF 'quantity' schema for scalars.""" 

199 

200 value: pydantic.StrictFloat 

201 unit: Unit 

202 

203 model_config = pydantic.ConfigDict( 

204 json_schema_extra={ 

205 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

206 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0", 

207 "tag": "!unit/quantity-1.2.0", 

208 } 

209 ) 

210 

211 

212class InlineArrayQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

213 """Model for a subset of the ASDF 'quantity' schema for inline arrays.""" 

214 

215 value: InlineArrayModel 

216 unit: Unit 

217 

218 model_config = pydantic.ConfigDict( 

219 json_schema_extra={ 

220 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

221 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0", 

222 "tag": "!unit/quantity-1.2.0", 

223 } 

224 ) 

225 

226 

227class ArrayReferenceQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

228 """Model for a subset of the ASDF 'quantity' schema for external arrays.""" 

229 

230 value: ArrayReferenceModel 

231 unit: Unit 

232 

233 model_config = pydantic.ConfigDict( 

234 json_schema_extra={ 

235 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

236 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0", 

237 "tag": "!unit/quantity-1.2.0", 

238 } 

239 ) 

240 

241 

242class _QuantitySerialization: 

243 """Pydantic hooks for scalar quantity serialization.""" 

244 

245 @classmethod 

246 def __get_pydantic_core_schema__( 

247 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

248 ) -> pcs.CoreSchema: 

249 from_model_schema = pcs.chain_schema( 

250 [ 

251 handler(QuantityModel), 

252 pcs.no_info_plain_validator_function(cls.from_model), 

253 ] 

254 ) 

255 return pcs.json_or_python_schema( 

256 json_schema=from_model_schema, 

257 python_schema=pcs.union_schema( 

258 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema] 

259 ), 

260 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model), 

261 ) 

262 

263 @classmethod 

264 def __get_pydantic_json_schema__( 

265 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler 

266 ) -> JsonSchemaValue: 

267 return handler(QuantityModel.__pydantic_core_schema__) 

268 

269 @classmethod 

270 def from_model(cls, model: QuantityModel) -> astropy.units.Quantity: 

271 return astropy.units.Quantity(model.value, unit=model.unit) 

272 

273 @classmethod 

274 def to_model(cls, quantity: astropy.units.Quantity) -> QuantityModel: 

275 assert quantity.isscalar 

276 return QuantityModel(value=quantity.to_value(), unit=quantity.unit) 

277 

278 

279type Quantity = Annotated[astropy.units.Quantity, _QuantitySerialization] 

280 

281 

282class _InlineArrayQuantitySerialization: 

283 """Pydantic hooks for inline array quantity serialization.""" 

284 

285 @classmethod 

286 def __get_pydantic_core_schema__( 

287 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

288 ) -> pcs.CoreSchema: 

289 from_model_schema = pcs.chain_schema( 

290 [ 

291 handler(InlineArrayQuantityModel), 

292 pcs.no_info_plain_validator_function(cls.from_model), 

293 ] 

294 ) 

295 return pcs.json_or_python_schema( 

296 json_schema=from_model_schema, 

297 python_schema=pcs.union_schema( 

298 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema] 

299 ), 

300 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model), 

301 ) 

302 

303 @classmethod 

304 def __get_pydantic_json_schema__( 

305 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler 

306 ) -> JsonSchemaValue: 

307 return handler(InlineArrayQuantityModel.__pydantic_core_schema__) 

308 

309 @classmethod 

310 def from_model(cls, model: InlineArrayQuantityModel) -> astropy.units.Quantity: 

311 return astropy.units.Quantity(_InlineArraySerialization.from_model(model.value), unit=model.unit) 

312 

313 @classmethod 

314 def to_model(cls, quantity: astropy.units.Quantity) -> InlineArrayQuantityModel: 

315 assert quantity.isscalar 

316 return InlineArrayQuantityModel( 

317 value=_InlineArraySerialization.to_model(quantity.to_value()), 

318 unit=quantity.unit, 

319 ) 

320 

321 

322type InlineArrayQuantity = Annotated[astropy.units.Quantity, _InlineArrayQuantitySerialization] 

323 

324 

325class TimeModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

326 """Model for a subset of the ASDF 'time' schema.""" 

327 

328 value: str 

329 scale: Literal["utc", "tai"] 

330 format: Literal["iso"] = "iso" 

331 

332 model_config = pydantic.ConfigDict( 

333 json_schema_extra={ 

334 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

335 "id": "http://stsci.edu/schemas/asdf/time/time-1.2.0", 

336 "tag": "!time/time-1.2.0", 

337 } 

338 ) 

339 

340 

341class _TimeSerialization: 

342 """Pydantic hooks for time serialization. 

343 

344 This class provides implementations for the `Time` type alias for 

345 `astropy.time.Time` that adds Pydantic serialization and validation. 

346 """ 

347 

348 @classmethod 

349 def __get_pydantic_core_schema__( 

350 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

351 ) -> pcs.CoreSchema: 

352 from_model_schema = pcs.chain_schema( 

353 [ 

354 TimeModel.__pydantic_core_schema__, 

355 pcs.no_info_plain_validator_function(cls.from_model), 

356 ] 

357 ) 

358 return pcs.json_or_python_schema( 

359 json_schema=from_model_schema, 

360 python_schema=pcs.union_schema([pcs.is_instance_schema(astropy.time.Time), from_model_schema]), 

361 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model, info_arg=False), 

362 ) 

363 

364 @classmethod 

365 def __get_pydantic_json_schema__( 

366 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler 

367 ) -> JsonSchemaValue: 

368 return handler(TimeModel.__pydantic_core_schema__) 

369 

370 @classmethod 

371 def from_model(cls, model: TimeModel) -> astropy.time.Time: 

372 return astropy.time.Time(model.value, scale=model.scale, format=model.format) 

373 

374 @classmethod 

375 def to_model(cls, time: astropy.time.Time) -> TimeModel: 

376 if time.scale != "utc" and time.scale != "tai": 

377 time = time.tai 

378 return TimeModel(value=time.to_value("iso"), scale=time.scale, format="iso") 

379 

380 

381type Time = Annotated[astropy.time.Time, _TimeSerialization]