Coverage for python/lsst/images/serialization/_asdf_utils.py: 68%
136 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-05-27 08:29 +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."""
110 return ArrayReferenceQuantityModel.model_construct(value=self, unit=unit)
112 model_config = pydantic.ConfigDict(
113 json_schema_extra={
114 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
115 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0",
116 "tag": "!core/ndarray-1.1.0",
117 }
118 )
121class InlineArrayModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
122 """Model for a subset of the ASDF 'ndarray' schema, in the case where the
123 array data is stored inline.
124 """
126 data: list[Any]
127 datatype: NumberType
129 @property
130 def shape(self) -> tuple[int, ...]:
131 """The shape of the array (`tuple` [`int`, ...])."""
132 return self._extract_shape(self.data)
134 def with_units(self, unit: astropy.unit.UnitBase) -> InlineArrayQuantityModel:
135 """Add units, transforming this model in to a Quantity model."""
136 return InlineArrayQuantityModel.model_construct(value=self, unit=unit)
138 @classmethod
139 def _extract_shape(cls, data: list[Any], current: tuple[int, ...] = ()) -> tuple[int, ...]:
140 if not data:
141 return current + (0,)
142 if not isinstance(data[0], list):
143 return current + (len(data),)
144 return cls._extract_shape(data[0], current + (len(data),))
146 model_config = pydantic.ConfigDict(
147 json_schema_extra={
148 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
149 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0",
150 "tag": "!core/ndarray-1.1.0",
151 }
152 )
155class _InlineArraySerialization:
156 """Pydantic hooks for array serialization.
158 This class provides implementations for the `Array` type alias for
159 `numpy.ndarray` that adds Pydantic serialization and validation.
160 """
162 @classmethod
163 def __get_pydantic_core_schema__(
164 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
165 ) -> pcs.CoreSchema:
166 from_model_schema = pcs.chain_schema(
167 [
168 handler(InlineArrayModel),
169 pcs.no_info_plain_validator_function(cls.from_model),
170 ]
171 )
172 return pcs.json_or_python_schema(
173 json_schema=from_model_schema,
174 python_schema=pcs.union_schema([pcs.is_instance_schema(np.ndarray), from_model_schema]),
175 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
176 )
178 @classmethod
179 def __get_pydantic_json_schema__(
180 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler
181 ) -> JsonSchemaValue:
182 return handler(InlineArrayModel.__pydantic_core_schema__)
184 @classmethod
185 def from_model(cls, model: InlineArrayModel) -> np.ndarray:
186 return np.array(model.data, dtype=model.datatype.to_numpy())
188 @classmethod
189 def to_model(cls, array: np.ndarray) -> InlineArrayModel:
190 datatype = NumberType.from_numpy(array.dtype)
191 return InlineArrayModel(data=array.tolist(), datatype=datatype)
194type InlineArray = Annotated[np.ndarray, _InlineArraySerialization]
197class QuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
198 """Model for a subset of the ASDF 'quantity' schema for scalars."""
200 value: pydantic.StrictFloat
201 unit: Unit
203 model_config = pydantic.ConfigDict(
204 json_schema_extra={
205 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
206 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
207 "tag": "!unit/quantity-1.2.0",
208 }
209 )
212class InlineArrayQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
213 """Model for a subset of the ASDF 'quantity' schema for inline arrays."""
215 value: InlineArrayModel
216 unit: Unit
218 model_config = pydantic.ConfigDict(
219 json_schema_extra={
220 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
221 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
222 "tag": "!unit/quantity-1.2.0",
223 }
224 )
227class ArrayReferenceQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
228 """Model for a subset of the ASDF 'quantity' schema for external arrays."""
230 value: ArrayReferenceModel
231 unit: Unit
233 model_config = pydantic.ConfigDict(
234 json_schema_extra={
235 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
236 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0",
237 "tag": "!unit/quantity-1.2.0",
238 }
239 )
242class _QuantitySerialization:
243 """Pydantic hooks for scalar quantity serialization."""
245 @classmethod
246 def __get_pydantic_core_schema__(
247 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
248 ) -> pcs.CoreSchema:
249 from_model_schema = pcs.chain_schema(
250 [
251 handler(QuantityModel),
252 pcs.no_info_plain_validator_function(cls.from_model),
253 ]
254 )
255 return pcs.json_or_python_schema(
256 json_schema=from_model_schema,
257 python_schema=pcs.union_schema(
258 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema]
259 ),
260 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
261 )
263 @classmethod
264 def __get_pydantic_json_schema__(
265 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler
266 ) -> JsonSchemaValue:
267 return handler(QuantityModel.__pydantic_core_schema__)
269 @classmethod
270 def from_model(cls, model: QuantityModel) -> astropy.units.Quantity:
271 return astropy.units.Quantity(model.value, unit=model.unit)
273 @classmethod
274 def to_model(cls, quantity: astropy.units.Quantity) -> QuantityModel:
275 assert quantity.isscalar
276 return QuantityModel(value=quantity.to_value(), unit=quantity.unit)
279type Quantity = Annotated[astropy.units.Quantity, _QuantitySerialization]
282class _InlineArrayQuantitySerialization:
283 """Pydantic hooks for inline array quantity serialization."""
285 @classmethod
286 def __get_pydantic_core_schema__(
287 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
288 ) -> pcs.CoreSchema:
289 from_model_schema = pcs.chain_schema(
290 [
291 handler(InlineArrayQuantityModel),
292 pcs.no_info_plain_validator_function(cls.from_model),
293 ]
294 )
295 return pcs.json_or_python_schema(
296 json_schema=from_model_schema,
297 python_schema=pcs.union_schema(
298 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema]
299 ),
300 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model),
301 )
303 @classmethod
304 def __get_pydantic_json_schema__(
305 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler
306 ) -> JsonSchemaValue:
307 return handler(InlineArrayQuantityModel.__pydantic_core_schema__)
309 @classmethod
310 def from_model(cls, model: InlineArrayQuantityModel) -> astropy.units.Quantity:
311 return astropy.units.Quantity(_InlineArraySerialization.from_model(model.value), unit=model.unit)
313 @classmethod
314 def to_model(cls, quantity: astropy.units.Quantity) -> InlineArrayQuantityModel:
315 assert quantity.isscalar
316 return InlineArrayQuantityModel(
317 value=_InlineArraySerialization.to_model(quantity.to_value()),
318 unit=quantity.unit,
319 )
322type InlineArrayQuantity = Annotated[astropy.units.Quantity, _InlineArrayQuantitySerialization]
325class TimeModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
326 """Model for a subset of the ASDF 'time' schema."""
328 value: str
329 scale: Literal["utc", "tai"]
330 format: Literal["iso"] = "iso"
332 model_config = pydantic.ConfigDict(
333 json_schema_extra={
334 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01",
335 "id": "http://stsci.edu/schemas/asdf/time/time-1.2.0",
336 "tag": "!time/time-1.2.0",
337 }
338 )
341class _TimeSerialization:
342 """Pydantic hooks for time serialization.
344 This class provides implementations for the `Time` type alias for
345 `astropy.time.Time` that adds Pydantic serialization and validation.
346 """
348 @classmethod
349 def __get_pydantic_core_schema__(
350 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
351 ) -> pcs.CoreSchema:
352 from_model_schema = pcs.chain_schema(
353 [
354 TimeModel.__pydantic_core_schema__,
355 pcs.no_info_plain_validator_function(cls.from_model),
356 ]
357 )
358 return pcs.json_or_python_schema(
359 json_schema=from_model_schema,
360 python_schema=pcs.union_schema([pcs.is_instance_schema(astropy.time.Time), from_model_schema]),
361 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model, info_arg=False),
362 )
364 @classmethod
365 def __get_pydantic_json_schema__(
366 cls, schema: pcs.CoreSchema, handler: GetJsonSchemaHandler
367 ) -> JsonSchemaValue:
368 return handler(TimeModel.__pydantic_core_schema__)
370 @classmethod
371 def from_model(cls, model: TimeModel) -> astropy.time.Time:
372 return astropy.time.Time(model.value, scale=model.scale, format=model.format)
374 @classmethod
375 def to_model(cls, time: astropy.time.Time) -> TimeModel:
376 if time.scale != "utc" and time.scale != "tai":
377 time = time.tai
378 return TimeModel(value=time.to_value("iso"), scale=time.scale, format="iso")
381type Time = Annotated[astropy.time.Time, _TimeSerialization]