Coverage for python/lsst/images/serialization/_asdf_utils.py: 88%
136 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-01 09:12 +0000
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.
12from __future__ import annotations
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)
28from typing import Annotated, Any, Literal
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
37from ._dtypes import NumberType
40class _UnitSerialization:
41 """Pydantic hooks for unit serialization.
43 This class provides implementations for the `Unit` type alias for
44 `astropy.unit.Unit` that adds Pydantic serialization and validation.
45 """
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 )
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)
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()
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]
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 """
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.")
108 def with_units(self, unit: astropy.units.UnitBase) -> ArrayReferenceQuantityModel:
109 """Add units, transforming this model into a Quantity model.
111 Parameters
112 ----------
113 unit
114 Units to attach to the array values.
115 """
116 return ArrayReferenceQuantityModel.model_construct(value=self, unit=unit)
118 model_config = pydantic.ConfigDict(
119 json_schema_extra={
120 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
121 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0",
122 "tag": "!core/ndarray-1.1.0",
123 }
124 )
127class InlineArrayModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
128 """Model for a subset of the ASDF 'ndarray' schema, in the case where the
129 array data is stored inline.
130 """
132 data: list[Any]
133 datatype: NumberType
135 @property
136 def shape(self) -> tuple[int, ...]:
137 """The shape of the array (`tuple` [`int`, ...])."""
138 return self._extract_shape(self.data)
140 def with_units(self, unit: astropy.unit.UnitBase) -> InlineArrayQuantityModel:
141 """Add units, transforming this model in to a Quantity model.
143 Parameters
144 ----------
145 unit
146 Units to attach to the array values.
147 """
148 return InlineArrayQuantityModel.model_construct(value=self, unit=unit)
150 @classmethod
151 def _extract_shape(cls, data: list[Any], current: tuple[int, ...] = ()) -> tuple[int, ...]:
152 if not data: 152 ↛ 153line 152 didn't jump to line 153 because the condition on line 152 was never true
153 return current + (0,)
154 if not isinstance(data[0], list):
155 return current + (len(data),)
156 return cls._extract_shape(data[0], current + (len(data),))
158 model_config = pydantic.ConfigDict(
159 json_schema_extra={
160 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
161 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0",
162 "tag": "!core/ndarray-1.1.0",
163 }
164 )
167class _InlineArraySerialization:
168 """Pydantic hooks for array serialization.
170 This class provides implementations for the `Array` type alias for
171 `numpy.ndarray` that adds Pydantic serialization and validation.
172 """
174 @classmethod
175 def __get_pydantic_core_schema__(
176 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
177 ) -> pcs.CoreSchema:
178 from_model_schema = pcs.chain_schema(
179 [
180 handler(InlineArrayModel),
181 pcs.no_info_plain_validator_function(cls.from_model),
182 ]
183 )
184 return pcs.json_or_python_schema(
185 json_schema=from_model_schema,
186 python_schema=pcs.union_schema([pcs.is_instance_schema(np.ndarray), from_model_schema]),
187 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
188 )
190 @classmethod
191 def __get_pydantic_json_schema__(
192 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler
193 ) -> JsonSchemaValue:
194 return handler(InlineArrayModel.__pydantic_core_schema__)
196 @classmethod
197 def from_model(cls, model: InlineArrayModel) -> np.ndarray:
198 return np.array(model.data, dtype=model.datatype.to_numpy())
200 @classmethod
201 def to_model(cls, array: np.ndarray) -> InlineArrayModel:
202 datatype = NumberType.from_numpy(array.dtype)
203 return InlineArrayModel(data=array.tolist(), datatype=datatype)
206type InlineArray = Annotated[np.ndarray, _InlineArraySerialization]
209class QuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
210 """Model for a subset of the ASDF 'quantity' schema for scalars."""
212 value: pydantic.StrictFloat
213 unit: Unit
215 model_config = pydantic.ConfigDict(
216 json_schema_extra={
217 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
218 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
219 "tag": "!unit/quantity-1.2.0",
220 }
221 )
224class InlineArrayQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
225 """Model for a subset of the ASDF 'quantity' schema for inline arrays."""
227 value: InlineArrayModel
228 unit: Unit
230 model_config = pydantic.ConfigDict(
231 json_schema_extra={
232 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
233 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
234 "tag": "!unit/quantity-1.2.0",
235 }
236 )
239class ArrayReferenceQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
240 """Model for a subset of the ASDF 'quantity' schema for external arrays."""
242 value: ArrayReferenceModel
243 unit: Unit
245 model_config = pydantic.ConfigDict(
246 json_schema_extra={
247 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
248 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
249 "tag": "!unit/quantity-1.2.0",
250 }
251 )
254class _QuantitySerialization:
255 """Pydantic hooks for scalar quantity serialization."""
257 @classmethod
258 def __get_pydantic_core_schema__(
259 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
260 ) -> pcs.CoreSchema:
261 from_model_schema = pcs.chain_schema(
262 [
263 handler(QuantityModel),
264 pcs.no_info_plain_validator_function(cls.from_model),
265 ]
266 )
267 return pcs.json_or_python_schema(
268 json_schema=from_model_schema,
269 python_schema=pcs.union_schema(
270 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema]
271 ),
272 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
273 )
275 @classmethod
276 def __get_pydantic_json_schema__(
277 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler
278 ) -> JsonSchemaValue:
279 return handler(QuantityModel.__pydantic_core_schema__)
281 @classmethod
282 def from_model(cls, model: QuantityModel) -> astropy.units.Quantity:
283 return astropy.units.Quantity(model.value, unit=model.unit)
285 @classmethod
286 def to_model(cls, quantity: astropy.units.Quantity) -> QuantityModel:
287 assert quantity.isscalar
288 return QuantityModel(value=quantity.to_value(), unit=quantity.unit)
291type Quantity = Annotated[astropy.units.Quantity, _QuantitySerialization]
294class _InlineArrayQuantitySerialization:
295 """Pydantic hooks for inline array quantity serialization."""
297 @classmethod
298 def __get_pydantic_core_schema__(
299 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
300 ) -> pcs.CoreSchema:
301 from_model_schema = pcs.chain_schema(
302 [
303 handler(InlineArrayQuantityModel),
304 pcs.no_info_plain_validator_function(cls.from_model),
305 ]
306 )
307 return pcs.json_or_python_schema(
308 json_schema=from_model_schema,
309 python_schema=pcs.union_schema(
310 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema]
311 ),
312 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
313 )
315 @classmethod
316 def __get_pydantic_json_schema__(
317 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler
318 ) -> JsonSchemaValue:
319 return handler(InlineArrayQuantityModel.__pydantic_core_schema__)
321 @classmethod
322 def from_model(cls, model: InlineArrayQuantityModel) -> astropy.units.Quantity:
323 return astropy.units.Quantity(_InlineArraySerialization.from_model(model.value), unit=model.unit)
325 @classmethod
326 def to_model(cls, quantity: astropy.units.Quantity) -> InlineArrayQuantityModel:
327 assert quantity.isscalar
328 return InlineArrayQuantityModel(
329 value=_InlineArraySerialization.to_model(quantity.to_value()),
330 unit=quantity.unit,
331 )
334type InlineArrayQuantity = Annotated[astropy.units.Quantity, _InlineArrayQuantitySerialization]
337class TimeModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
338 """Model for a subset of the ASDF 'time' schema."""
340 value: str
341 scale: Literal["utc", "tai"]
342 format: Literal["iso"] = "iso"
344 model_config = pydantic.ConfigDict(
345 json_schema_extra={
346 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
347 "id": "http://stsci.edu/schemas/asdf/time/time-1.2.0",
348 "tag": "!time/time-1.2.0",
349 }
350 )
353class _TimeSerialization:
354 """Pydantic hooks for time serialization.
356 This class provides implementations for the `Time` type alias for
357 `astropy.time.Time` that adds Pydantic serialization and validation.
358 """
360 @classmethod
361 def __get_pydantic_core_schema__(
362 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
363 ) -> pcs.CoreSchema:
364 from_model_schema = pcs.chain_schema(
365 [
366 TimeModel.__pydantic_core_schema__,
367 pcs.no_info_plain_validator_function(cls.from_model),
368 ]
369 )
370 return pcs.json_or_python_schema(
371 json_schema=from_model_schema,
372 python_schema=pcs.union_schema([pcs.is_instance_schema(astropy.time.Time), from_model_schema]),
373 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model, info_arg=False),
374 )
376 @classmethod
377 def __get_pydantic_json_schema__(
378 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler
379 ) -> JsonSchemaValue:
380 return handler(TimeModel.__pydantic_core_schema__)
382 @classmethod
383 def from_model(cls, model: TimeModel) -> astropy.time.Time:
384 return astropy.time.Time(model.value, scale=model.scale, format=model.format)
386 @classmethod
387 def to_model(cls, time: astropy.time.Time) -> TimeModel:
388 if time.scale != "utc" and time.scale != "tai":
389 time = time.tai
390 return TimeModel(value=time.to_value("iso"), scale=time.scale, format="iso")
393type Time = Annotated[astropy.time.Time, _TimeSerialization]