Coverage for python / lsst / images / serialization / _asdf_utils.py: 68%
124 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-13 08:46 +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 try:
67 return astropy.units.Unit(value, format="vounit")
68 except ValueError:
69 pass
70 # Some important units (e.g. "dn") are not supported by vounit, so
71 # fall back to letting astropy to infer the format.
72 return astropy.units.Unit(value)
74 @staticmethod
75 def to_str(unit: astropy.units.UnitBase) -> str:
76 try:
77 return unit.to_string("vounit")
78 except ValueError:
79 pass
80 # Some important units (e.g. "dn") are not supported by vounit, so
81 # fall back to letting astropy use the default format.
82 return unit.to_string()
85type Unit = Annotated[
86 astropy.units.UnitBase,
87 _UnitSerialization,
88 pydantic.WithJsonSchema(
89 {
90 "type": "string",
91 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
92 "id": "http://stsci.edu/schemas/asdf/unit/unit-1.0.0",
93 "tag": "!unit/unit-1.0.0",
94 }
95 ),
96]
99class ArrayReferenceModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
100 """Model for a subset of the ASDF 'ndarray' schema, in the case where the
101 array data is stored elsewhere.
102 """
104 source: str | int = pydantic.Field(description="Location of the underlying binary data.")
105 shape: list[int] = pydantic.Field(
106 # In (e.g.) FITS this is stored outside of the JSON as well, and it
107 # be hard to get it right if we need to make a reference to a column
108 # before all rows have been written, so unlike ASDF we allow this to
109 # be omitted.
110 default_factory=list,
111 description="Size of the array in each dimension.",
112 exclude_if=operator.not_,
113 )
114 datatype: NumberType = pydantic.Field(description="Data type of the array.")
115 byteorder: Literal["big"] = pydantic.Field(default="big", description="Byte order for the binary data.")
117 def with_units(self, unit: astropy.units.UnitBase) -> ArrayReferenceQuantityModel:
118 """Add units, transforming this model into a Quantity model."""
119 return ArrayReferenceQuantityModel.model_construct(value=self, unit=unit)
121 model_config = pydantic.ConfigDict(
122 json_schema_extra={
123 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
124 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0",
125 "tag": "!core/ndarray-1.1.0",
126 }
127 )
130class InlineArrayModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
131 """Model for a subset of the ASDF 'ndarray' schema, in the case where the
132 array data is stored inline.
133 """
135 data: list[Any]
136 datatype: NumberType
138 @property
139 def shape(self) -> tuple[int, ...]:
140 """The shape of the array (`tuple` [`int`, ...])."""
141 return self._extract_shape(self.data)
143 def with_units(self, unit: astropy.unit.UnitBase) -> InlineArrayQuantityModel:
144 """Add units, transforming this model in to a Quantity model."""
145 return InlineArrayQuantityModel.model_construct(value=self, unit=unit)
147 @classmethod
148 def _extract_shape(cls, data: list[Any], current: tuple[int, ...] = ()) -> tuple[int, ...]:
149 if not data:
150 return current + (0,)
151 if not isinstance(data[0], list):
152 return current + (len(data),)
153 return cls._extract_shape(data[0], current + (len(data),))
155 model_config = pydantic.ConfigDict(
156 json_schema_extra={
157 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
158 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0",
159 "tag": "!core/ndarray-1.1.0",
160 }
161 )
164class _InlineArraySerialization:
165 """Pydantic hooks for array serialization.
167 This class provides implementations for the `Array` type alias for
168 `numpy.ndarray` that adds Pydantic serialization and validation.
169 """
171 @classmethod
172 def __get_pydantic_core_schema__(
173 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
174 ) -> pcs.CoreSchema:
175 from_model_schema = pcs.chain_schema(
176 [
177 handler(InlineArrayModel),
178 pcs.no_info_plain_validator_function(cls.from_model),
179 ]
180 )
181 return pcs.json_or_python_schema(
182 json_schema=from_model_schema,
183 python_schema=pcs.union_schema([pcs.is_instance_schema(np.ndarray), from_model_schema]),
184 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
185 )
187 @classmethod
188 def from_model(cls, model: InlineArrayModel) -> np.ndarray:
189 return np.array(model.data, dtype=model.datatype.to_numpy())
191 @classmethod
192 def to_model(cls, array: np.ndarray) -> InlineArrayModel:
193 datatype = NumberType.from_numpy(array.dtype)
194 return InlineArrayModel(data=array.tolist(), datatype=datatype)
197type InlineArray = Annotated[np.ndarray, _InlineArraySerialization]
200class QuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
201 """Model for a subset of the ASDF 'quantity' schema for scalars."""
203 value: pydantic.StrictFloat
204 unit: Unit
206 model_config = pydantic.ConfigDict(
207 json_schema_extra={
208 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
209 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
210 "tag": "!unit/quantity-1.2.0",
211 }
212 )
215class InlineArrayQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
216 """Model for a subset of the ASDF 'quantity' schema for inline arrays."""
218 value: InlineArrayModel
219 unit: Unit
221 model_config = pydantic.ConfigDict(
222 json_schema_extra={
223 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
224 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
225 "tag": "!unit/quantity-1.2.0",
226 }
227 )
230class ArrayReferenceQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
231 """Model for a subset of the ASDF 'quantity' schema for external arrays."""
233 value: ArrayReferenceModel
234 unit: Unit
236 model_config = pydantic.ConfigDict(
237 json_schema_extra={
238 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
239 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
240 "tag": "!unit/quantity-1.2.0",
241 }
242 )
245class _QuantitySerialization:
246 """Pydantic hooks for scalar quantity serialization."""
248 @classmethod
249 def __get_pydantic_core_schema__(
250 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
251 ) -> pcs.CoreSchema:
252 from_model_schema = pcs.chain_schema(
253 [
254 handler(QuantityModel),
255 pcs.no_info_plain_validator_function(cls.from_model),
256 ]
257 )
258 return pcs.json_or_python_schema(
259 json_schema=from_model_schema,
260 python_schema=pcs.union_schema(
261 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema]
262 ),
263 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
264 )
266 @classmethod
267 def from_model(cls, model: QuantityModel) -> astropy.units.Quantity:
268 return astropy.units.Quantity(model.value, unit=model.unit)
270 @classmethod
271 def to_model(cls, quantity: astropy.units.Quantity) -> QuantityModel:
272 assert quantity.isscalar
273 return QuantityModel(value=quantity.to_value(), unit=quantity.unit)
276type Quantity = Annotated[astropy.units.Quantity, _QuantitySerialization]
279class _InlineArrayQuantitySerialization:
280 """Pydantic hooks for inline array quantity serialization."""
282 @classmethod
283 def __get_pydantic_core_schema__(
284 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
285 ) -> pcs.CoreSchema:
286 from_model_schema = pcs.chain_schema(
287 [
288 handler(InlineArrayQuantityModel),
289 pcs.no_info_plain_validator_function(cls.from_model),
290 ]
291 )
292 return pcs.json_or_python_schema(
293 json_schema=from_model_schema,
294 python_schema=pcs.union_schema(
295 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema]
296 ),
297 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
298 )
300 @classmethod
301 def from_model(cls, model: InlineArrayQuantityModel) -> astropy.units.Quantity:
302 return astropy.units.Quantity(_InlineArraySerialization.from_model(model.value), unit=model.unit)
304 @classmethod
305 def to_model(cls, quantity: astropy.units.Quantity) -> InlineArrayQuantityModel:
306 assert quantity.isscalar
307 return InlineArrayQuantityModel(
308 value=_InlineArraySerialization.to_model(quantity.to_value()),
309 unit=quantity.unit,
310 )
313type InlineArrayQuantity = Annotated[astropy.units.Quantity, _InlineArrayQuantitySerialization]
316class TimeModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
317 """Model for a subset of the ASDF 'time' schema."""
319 value: str
320 scale: Literal["utc", "tai"]
321 format: Literal["iso"] = "iso"
323 model_config = pydantic.ConfigDict(
324 json_schema_extra={
325 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
326 "id": "http://stsci.edu/schemas/asdf/time/time-1.2.0",
327 "tag": "!time/time-1.2.0",
328 }
329 )
332class _TimeSerialization:
333 """Pydantic hooks for time serialization.
335 This class provides implementations for the `Time` type alias for
336 `astropy.time.Time` that adds Pydantic serialization and validation.
337 """
339 @classmethod
340 def __get_pydantic_core_schema__(
341 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
342 ) -> pcs.CoreSchema:
343 from_model_schema = pcs.chain_schema(
344 [
345 TimeModel.__pydantic_core_schema__,
346 pcs.no_info_plain_validator_function(cls.from_model),
347 ]
348 )
349 return pcs.json_or_python_schema(
350 json_schema=from_model_schema,
351 python_schema=pcs.union_schema([pcs.is_instance_schema(astropy.time.Time), from_model_schema]),
352 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model, info_arg=False),
353 )
355 @classmethod
356 def from_model(cls, model: TimeModel) -> astropy.time.Time:
357 return astropy.time.Time(model.value, scale=model.scale, format=model.format)
359 @classmethod
360 def to_model(cls, time: astropy.time.Time) -> TimeModel:
361 if time.scale != "utc" and time.scale != "tai":
362 time = time.tai
363 return TimeModel(value=time.to_value("iso"), scale=time.scale, format="iso")
366type Time = Annotated[astropy.time.Time, _TimeSerialization]