Coverage for python / lsst / images / serialization / _asdf_utils.py: 71%
116 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-30 09:07 +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
14import operator
16__all__ = (
17 "ArrayReferenceModel",
18 "ArrayReferenceQuantityModel",
19 "InlineArray",
20 "InlineArrayModel",
21 "InlineArrayQuantity",
22 "InlineArrayQuantityModel",
23 "Quantity",
24 "QuantityModel",
25 "Time",
26 "TimeModel",
27 "Unit",
28)
30from typing import Annotated, Any, Literal
32import astropy.time
33import astropy.units
34import numpy as np
35import pydantic
36import pydantic_core.core_schema as pcs
38from ._dtypes import NumberType
41class _UnitSerialization:
42 """Pydantic hooks for unit serialization.
44 This class provides implementations for the `Unit` type alias for
45 `astropy.unit.Unit` that adds Pydantic serialization and validation.
46 """
48 @classmethod
49 def __get_pydantic_core_schema__(
50 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
51 ) -> pcs.CoreSchema:
52 from_str_schema = pcs.chain_schema(
53 [
54 pcs.str_schema(),
55 pcs.no_info_plain_validator_function(cls.from_str),
56 ]
57 )
58 return pcs.json_or_python_schema(
59 json_schema=from_str_schema,
60 python_schema=pcs.union_schema([pcs.is_instance_schema(astropy.units.UnitBase), from_str_schema]),
61 serialization=pcs.plain_serializer_function_ser_schema(cls.to_str),
62 )
64 @classmethod
65 def from_str(cls, value: str) -> astropy.units.UnitBase:
66 return astropy.units.Unit(value, format="vounit")
68 @staticmethod
69 def to_str(unit: astropy.units.UnitBase) -> str:
70 return unit.to_string("vounit")
73type Unit = Annotated[
74 astropy.units.UnitBase,
75 _UnitSerialization,
76 pydantic.WithJsonSchema(
77 {
78 "type": "string",
79 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
80 "id": "http://stsci.edu/schemas/asdf/unit/unit-1.0.0",
81 "tag": "!unit/unit-1.0.0",
82 }
83 ),
84]
87class ArrayReferenceModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
88 """Model for a subset of the ASDF 'ndarray' schema, in the case where the
89 array data is stored elsewhere.
90 """
92 source: str | int = pydantic.Field(description="Location of the underlying binary data.")
93 shape: list[int] = pydantic.Field(
94 # In (e.g.) FITS this is stored outside of the JSON as well, and it
95 # be hard to get it right if we need to make a reference to a column
96 # before all rows have been written, so unlike ASDF we allow this to
97 # be omitted.
98 default_factory=list,
99 description="Size of the array in each dimension.",
100 exclude_if=operator.not_,
101 )
102 datatype: NumberType = pydantic.Field(description="Data type of the array.")
103 byteorder: Literal["big"] = pydantic.Field(default="big", description="Byte order for the binary data.")
105 def with_units(self, unit: astropy.units.UnitBase) -> ArrayReferenceQuantityModel:
106 """Add units, transforming this model into a Quantity model."""
107 return ArrayReferenceQuantityModel.model_construct(value=self, unit=unit)
109 model_config = pydantic.ConfigDict(
110 json_schema_extra={
111 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
112 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0",
113 "tag": "!core/ndarray-1.1.0",
114 }
115 )
118class InlineArrayModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
119 """Model for a subset of the ASDF 'ndarray' schema, in the case where the
120 array data is stored inline.
121 """
123 data: list[Any]
124 datatype: NumberType
126 @property
127 def shape(self) -> tuple[int, ...]:
128 """The shape of the array (`tuple` [`int`, ...])."""
129 return self._extract_shape(self.data)
131 def with_units(self, unit: astropy.unit.UnitBase) -> InlineArrayQuantityModel:
132 """Add units, transforming this model in to a Quantity model."""
133 return InlineArrayQuantityModel.model_construct(value=self, unit=unit)
135 @classmethod
136 def _extract_shape(cls, data: list[Any], current: tuple[int, ...] = ()) -> tuple[int, ...]:
137 if not data:
138 return current + (0,)
139 if not isinstance(data[0], list):
140 return current + (len(data),)
141 return cls._extract_shape(data[0], current + (len(data),))
143 model_config = pydantic.ConfigDict(
144 json_schema_extra={
145 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
146 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0",
147 "tag": "!core/ndarray-1.1.0",
148 }
149 )
152class _InlineArraySerialization:
153 """Pydantic hooks for array serialization.
155 This class provides implementations for the `Array` type alias for
156 `numpy.ndarray` that adds Pydantic serialization and validation.
157 """
159 @classmethod
160 def __get_pydantic_core_schema__(
161 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
162 ) -> pcs.CoreSchema:
163 from_model_schema = pcs.chain_schema(
164 [
165 handler(InlineArrayModel),
166 pcs.no_info_plain_validator_function(cls.from_model),
167 ]
168 )
169 return pcs.json_or_python_schema(
170 json_schema=from_model_schema,
171 python_schema=pcs.union_schema([pcs.is_instance_schema(np.ndarray), from_model_schema]),
172 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
173 )
175 @classmethod
176 def from_model(cls, model: InlineArrayModel) -> np.ndarray:
177 return np.array(model.data, dtype=model.datatype.to_numpy())
179 @classmethod
180 def to_model(cls, array: np.ndarray) -> InlineArrayModel:
181 datatype = NumberType.from_numpy(array.dtype)
182 return InlineArrayModel(data=array.tolist(), datatype=datatype)
185type InlineArray = Annotated[np.ndarray, _InlineArraySerialization]
188class QuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
189 """Model for a subset of the ASDF 'quantity' schema for scalars."""
191 value: pydantic.StrictFloat
192 unit: Unit
194 model_config = pydantic.ConfigDict(
195 json_schema_extra={
196 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
197 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
198 "tag": "!unit/quantity-1.2.0",
199 }
200 )
203class InlineArrayQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
204 """Model for a subset of the ASDF 'quantity' schema for inline arrays."""
206 value: InlineArrayModel
207 unit: Unit
209 model_config = pydantic.ConfigDict(
210 json_schema_extra={
211 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
212 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
213 "tag": "!unit/quantity-1.2.0",
214 }
215 )
218class ArrayReferenceQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
219 """Model for a subset of the ASDF 'quantity' schema for external arrays."""
221 value: ArrayReferenceModel
222 unit: Unit
224 model_config = pydantic.ConfigDict(
225 json_schema_extra={
226 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
227 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
228 "tag": "!unit/quantity-1.2.0",
229 }
230 )
233class _QuantitySerialization:
234 """Pydantic hooks for scalar quantity serialization."""
236 @classmethod
237 def __get_pydantic_core_schema__(
238 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
239 ) -> pcs.CoreSchema:
240 from_model_schema = pcs.chain_schema(
241 [
242 handler(QuantityModel),
243 pcs.no_info_plain_validator_function(cls.from_model),
244 ]
245 )
246 return pcs.json_or_python_schema(
247 json_schema=from_model_schema,
248 python_schema=pcs.union_schema(
249 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema]
250 ),
251 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
252 )
254 @classmethod
255 def from_model(cls, model: QuantityModel) -> astropy.units.Quantity:
256 return astropy.units.Quantity(model.value, unit=model.unit)
258 @classmethod
259 def to_model(cls, quantity: astropy.units.Quantity) -> QuantityModel:
260 assert quantity.isscalar
261 return QuantityModel(value=quantity.to_value(), unit=quantity.unit)
264type Quantity = Annotated[astropy.units.Quantity, _QuantitySerialization]
267class _InlineArrayQuantitySerialization:
268 """Pydantic hooks for inline array quantity serialization."""
270 @classmethod
271 def __get_pydantic_core_schema__(
272 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
273 ) -> pcs.CoreSchema:
274 from_model_schema = pcs.chain_schema(
275 [
276 handler(InlineArrayQuantityModel),
277 pcs.no_info_plain_validator_function(cls.from_model),
278 ]
279 )
280 return pcs.json_or_python_schema(
281 json_schema=from_model_schema,
282 python_schema=pcs.union_schema(
283 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema]
284 ),
285 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
286 )
288 @classmethod
289 def from_model(cls, model: InlineArrayQuantityModel) -> astropy.units.Quantity:
290 return astropy.units.Quantity(_InlineArraySerialization.from_model(model.value), unit=model.unit)
292 @classmethod
293 def to_model(cls, quantity: astropy.units.Quantity) -> InlineArrayQuantityModel:
294 assert quantity.isscalar
295 return InlineArrayQuantityModel(
296 value=_InlineArraySerialization.to_model(quantity.to_value()),
297 unit=quantity.unit,
298 )
301type InlineArrayQuantity = Annotated[astropy.units.Quantity, _InlineArrayQuantitySerialization]
304class TimeModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
305 """Model for a subset of the ASDF 'time' schema."""
307 value: str
308 scale: Literal["utc", "tai"]
309 format: Literal["iso"] = "iso"
311 model_config = pydantic.ConfigDict(
312 json_schema_extra={
313 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
314 "id": "http://stsci.edu/schemas/asdf/time/time-1.2.0",
315 "tag": "!time/time-1.2.0",
316 }
317 )
320class _TimeSerialization:
321 """Pydantic hooks for time serialization.
323 This class provides implementations for the `Time` type alias for
324 `astropy.time.Time` that adds Pydantic serialization and validation.
325 """
327 @classmethod
328 def __get_pydantic_core_schema__(
329 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
330 ) -> pcs.CoreSchema:
331 from_model_schema = pcs.chain_schema(
332 [
333 TimeModel.__pydantic_core_schema__,
334 pcs.no_info_plain_validator_function(cls.from_model),
335 ]
336 )
337 return pcs.json_or_python_schema(
338 json_schema=from_model_schema,
339 python_schema=pcs.union_schema([pcs.is_instance_schema(astropy.time.Time), from_model_schema]),
340 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model, info_arg=False),
341 )
343 @classmethod
344 def from_model(cls, model: TimeModel) -> astropy.time.Time:
345 return astropy.time.Time(model.value, scale=model.scale, format=model.format)
347 @classmethod
348 def to_model(cls, time: astropy.time.Time) -> TimeModel:
349 if time.scale != "utc" and time.scale != "tai":
350 time = time.tai
351 return TimeModel(value=time.to_value("iso"), scale=time.scale, format="iso")
354type Time = Annotated[astropy.time.Time, _TimeSerialization]