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

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, ClassVar, Literal 

29 

30import astropy.time 

31import astropy.units 

32import numpy as np 

33import pydantic 

34import pydantic_core.core_schema as pcs 

35 

36from ._dtypes import NumberType 

37 

38 

39class _UnitSerialization: 

40 """Pydantic hooks for unit serialization. 

41 

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

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

44 """ 

45 

46 @classmethod 

47 def __get_pydantic_core_schema__( 

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

49 ) -> pcs.CoreSchema: 

50 from_str_schema = pcs.chain_schema( 

51 [ 

52 pcs.str_schema(), 

53 pcs.no_info_plain_validator_function(cls.from_str), 

54 ] 

55 ) 

56 return pcs.json_or_python_schema( 

57 json_schema=from_str_schema, 

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

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

60 ) 

61 

62 @classmethod 

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

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

65 

66 @staticmethod 

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

68 return unit.to_string("vounit") 

69 

70 

71type Unit = Annotated[ 

72 astropy.units.UnitBase, 

73 _UnitSerialization, 

74 pydantic.WithJsonSchema( 

75 { 

76 "type": "string", 

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

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

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

80 } 

81 ), 

82] 

83 

84 

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

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

87 array data is stored elsewhere. 

88 """ 

89 

90 source: str | int 

91 shape: list[int] 

92 datatype: NumberType 

93 byteorder: Literal["big"] = "big" 

94 

95 model_config = pydantic.ConfigDict( 

96 json_schema_extra={ 

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

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

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

100 } 

101 ) 

102 

103 source_is_table: ClassVar[Literal[False]] = False 

104 

105 

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

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

108 array data is stored inline. 

109 """ 

110 

111 data: list[Any] 

112 datatype: NumberType 

113 

114 model_config = pydantic.ConfigDict( 

115 json_schema_extra={ 

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

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

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

119 } 

120 ) 

121 

122 

123class _InlineArraySerialization: 

124 """Pydantic hooks for array serialization. 

125 

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

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

128 """ 

129 

130 @classmethod 

131 def __get_pydantic_core_schema__( 

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

133 ) -> pcs.CoreSchema: 

134 from_model_schema = pcs.chain_schema( 

135 [ 

136 handler(InlineArrayModel), 

137 pcs.no_info_plain_validator_function(cls.from_model), 

138 ] 

139 ) 

140 return pcs.json_or_python_schema( 

141 json_schema=from_model_schema, 

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

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

144 ) 

145 

146 @classmethod 

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

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

149 

150 @classmethod 

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

152 datatype = NumberType.from_numpy(array.dtype) 

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

154 

155 

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

157 

158 

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

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

161 

162 value: pydantic.StrictFloat 

163 unit: Unit 

164 

165 model_config = pydantic.ConfigDict( 

166 json_schema_extra={ 

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

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

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

170 } 

171 ) 

172 

173 

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

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

176 

177 value: InlineArrayModel 

178 unit: Unit 

179 

180 model_config = pydantic.ConfigDict( 

181 json_schema_extra={ 

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

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

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

185 } 

186 ) 

187 

188 

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

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

191 

192 value: ArrayReferenceModel 

193 unit: Unit 

194 

195 model_config = pydantic.ConfigDict( 

196 json_schema_extra={ 

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

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

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

200 } 

201 ) 

202 

203 

204class _QuantitySerialization: 

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

206 

207 @classmethod 

208 def __get_pydantic_core_schema__( 

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

210 ) -> pcs.CoreSchema: 

211 from_model_schema = pcs.chain_schema( 

212 [ 

213 handler(QuantityModel), 

214 pcs.no_info_plain_validator_function(cls.from_model), 

215 ] 

216 ) 

217 return pcs.json_or_python_schema( 

218 json_schema=from_model_schema, 

219 python_schema=pcs.union_schema( 

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

221 ), 

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

223 ) 

224 

225 @classmethod 

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

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

228 

229 @classmethod 

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

231 assert quantity.isscalar 

232 return QuantityModel(value=quantity.to_value(), unit=_UnitSerialization.to_str(quantity.unit)) 

233 

234 

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

236 

237 

238class _InlineArrayQuantitySerialization: 

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

240 

241 @classmethod 

242 def __get_pydantic_core_schema__( 

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

244 ) -> pcs.CoreSchema: 

245 from_model_schema = pcs.chain_schema( 

246 [ 

247 handler(InlineArrayQuantityModel), 

248 pcs.no_info_plain_validator_function(cls.from_model), 

249 ] 

250 ) 

251 return pcs.json_or_python_schema( 

252 json_schema=from_model_schema, 

253 python_schema=pcs.union_schema( 

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

255 ), 

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

257 ) 

258 

259 @classmethod 

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

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

262 

263 @classmethod 

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

265 assert quantity.isscalar 

266 return InlineArrayQuantityModel( 

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

268 unit=_UnitSerialization.to_str(quantity.unit), 

269 ) 

270 

271 

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

273 

274 

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

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

277 

278 value: str 

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

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

281 

282 model_config = pydantic.ConfigDict( 

283 json_schema_extra={ 

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

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

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

287 } 

288 ) 

289 

290 

291class _TimeSerialization: 

292 """Pydantic hooks for time serialization. 

293 

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

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

296 """ 

297 

298 @classmethod 

299 def __get_pydantic_core_schema__( 

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

301 ) -> pcs.CoreSchema: 

302 from_model_schema = pcs.chain_schema( 

303 [ 

304 TimeModel.__pydantic_core_schema__, 

305 pcs.no_info_plain_validator_function(cls.from_model), 

306 ] 

307 ) 

308 return pcs.json_or_python_schema( 

309 json_schema=from_model_schema, 

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

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

312 ) 

313 

314 @classmethod 

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

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

317 

318 @classmethod 

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

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

321 time = time.tai 

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

323 

324 

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