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 ._common import ArchiveTree 

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 return astropy.units.Unit(value, format="vounit") 

66 

67 @staticmethod 

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

69 return unit.to_string("vounit") 

70 

71 

72type Unit = Annotated[ 

73 astropy.units.UnitBase, 

74 _UnitSerialization, 

75 pydantic.WithJsonSchema( 

76 { 

77 "type": "string", 

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

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

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

81 } 

82 ), 

83] 

84 

85 

86class ArrayReferenceModel(ArchiveTree): 

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

88 array data is stored elsewhere. 

89 """ 

90 

91 source: str | int 

92 shape: list[int] 

93 datatype: NumberType 

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

95 

96 model_config = pydantic.ConfigDict( 

97 json_schema_extra={ 

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

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

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

101 } 

102 ) 

103 

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

105 

106 

107class InlineArrayModel(ArchiveTree): 

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

109 array data is stored inline. 

110 """ 

111 

112 data: list[Any] 

113 datatype: NumberType 

114 

115 model_config = pydantic.ConfigDict( 

116 json_schema_extra={ 

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

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

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

120 } 

121 ) 

122 

123 

124class _InlineArraySerialization: 

125 """Pydantic hooks for array serialization. 

126 

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

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

129 """ 

130 

131 @classmethod 

132 def __get_pydantic_core_schema__( 

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

134 ) -> pcs.CoreSchema: 

135 from_model_schema = pcs.chain_schema( 

136 [ 

137 handler(InlineArrayModel), 

138 pcs.no_info_plain_validator_function(cls.from_model), 

139 ] 

140 ) 

141 return pcs.json_or_python_schema( 

142 json_schema=from_model_schema, 

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

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

145 ) 

146 

147 @classmethod 

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

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

150 

151 @classmethod 

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

153 datatype = NumberType.from_numpy(array.dtype) 

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

155 

156 

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

158 

159 

160class QuantityModel(ArchiveTree): 

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

162 

163 value: pydantic.StrictFloat 

164 unit: Unit 

165 

166 model_config = pydantic.ConfigDict( 

167 json_schema_extra={ 

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

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

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

171 } 

172 ) 

173 

174 

175class InlineArrayQuantityModel(ArchiveTree): 

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

177 

178 value: InlineArrayModel 

179 unit: Unit 

180 

181 model_config = pydantic.ConfigDict( 

182 json_schema_extra={ 

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

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

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

186 } 

187 ) 

188 

189 

190class ArrayReferenceQuantityModel(ArchiveTree): 

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

192 

193 value: ArrayReferenceModel 

194 unit: Unit 

195 

196 model_config = pydantic.ConfigDict( 

197 json_schema_extra={ 

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

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

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

201 } 

202 ) 

203 

204 

205class _QuantitySerialization: 

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

207 

208 @classmethod 

209 def __get_pydantic_core_schema__( 

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

211 ) -> pcs.CoreSchema: 

212 from_model_schema = pcs.chain_schema( 

213 [ 

214 handler(QuantityModel), 

215 pcs.no_info_plain_validator_function(cls.from_model), 

216 ] 

217 ) 

218 return pcs.json_or_python_schema( 

219 json_schema=from_model_schema, 

220 python_schema=pcs.union_schema( 

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

222 ), 

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

224 ) 

225 

226 @classmethod 

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

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

229 

230 @classmethod 

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

232 assert quantity.isscalar 

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

234 

235 

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

237 

238 

239class _InlineArrayQuantitySerialization: 

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

241 

242 @classmethod 

243 def __get_pydantic_core_schema__( 

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

245 ) -> pcs.CoreSchema: 

246 from_model_schema = pcs.chain_schema( 

247 [ 

248 handler(InlineArrayQuantityModel), 

249 pcs.no_info_plain_validator_function(cls.from_model), 

250 ] 

251 ) 

252 return pcs.json_or_python_schema( 

253 json_schema=from_model_schema, 

254 python_schema=pcs.union_schema( 

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

256 ), 

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

258 ) 

259 

260 @classmethod 

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

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

263 

264 @classmethod 

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

266 assert quantity.isscalar 

267 return InlineArrayQuantityModel( 

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

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

270 ) 

271 

272 

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

274 

275 

276class TimeModel(ArchiveTree): 

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

278 

279 value: str 

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

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

282 

283 model_config = pydantic.ConfigDict( 

284 json_schema_extra={ 

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

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

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

288 } 

289 ) 

290 

291 

292class _TimeSerialization: 

293 """Pydantic hooks for time serialization. 

294 

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

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

297 """ 

298 

299 @classmethod 

300 def __get_pydantic_core_schema__( 

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

302 ) -> pcs.CoreSchema: 

303 from_model_schema = pcs.chain_schema( 

304 [ 

305 TimeModel.__pydantic_core_schema__, 

306 pcs.no_info_plain_validator_function(cls.from_model), 

307 ] 

308 ) 

309 return pcs.json_or_python_schema( 

310 json_schema=from_model_schema, 

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

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

313 ) 

314 

315 @classmethod 

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

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

318 

319 @classmethod 

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

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

322 time = time.tai 

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

324 

325 

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