Coverage for python/lsst/images/serialization/_output_archive.py: 91%
57 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-14 20:37 -0700
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 "NestedOutputArchive",
16 "OutputArchive",
17)
19from abc import ABC, abstractmethod
20from collections.abc import Callable, Hashable, Iterator, Mapping
21from typing import TYPE_CHECKING, TypeVar
23import astropy.io.fits
24import astropy.table
25import astropy.units
26import numpy as np
27import pydantic
29from ._asdf_utils import ArrayReferenceModel, InlineArrayModel
30from ._common import ArchiveTree, no_header_updates
31from ._tables import TableModel
33if TYPE_CHECKING:
34 from .._transforms import FrameSet
36# This pre-python-3.12 declaration is needed by Sphinx (probably the
37# autodoc-typehints plugin.
38P = TypeVar("P", bound=pydantic.BaseModel)
41class OutputArchive[P](ABC):
42 """Abstract interface for writing to a file format.
44 Notes
45 -----
46 An output archive instance is assumed to be paired with a Pydantic model
47 that represents a JSON tree, with the archive used to serialize data that
48 is not natively JSON into data that is (which may just be a reference to
49 binary data stored elsewhere in the file). The archive doesn't actually
50 hold that model instance because we don't want to assume it can be built
51 via default-initialization and assignment, and because we'd prefer to avoid
52 making the output archive generic over the model type. It is expected that
53 most concrete archive implementations will accept the paired model in some
54 sort of finalization method in order to write it into the file, but this is
55 not part of the base class interface.
56 """
58 def __init__(self) -> None:
59 self._name_versions: dict[str, int] = {}
60 """Per-name occurrence count, used by `_register_name` to disambiguate
61 repeated logical names within a single write (e.g. each operand of a
62 `SumField` calling ``add_array(name="data")`` from the same nested
63 archive).
64 """
66 def _register_name(self, name: str) -> tuple[str, int]:
67 """Return the input name and its 1-based occurrence count.
69 Parameters
70 ----------
71 name
72 The logical archive name being saved (typically the absolute
73 archive path of an array, table, or pointer target).
75 Returns
76 -------
77 name : `str`
78 The input name, returned unchanged so that the caller controls
79 how the version is rendered into the on-disk identifier.
80 version : `int`
81 ``1`` the first time a given name is registered, then ``2``,
82 ``3`` and so on for subsequent calls with the same name.
84 Notes
85 -----
86 Backends should call this from `add_array`, `add_table`,
87 `add_structured_array`, and `serialize_pointer` to detect repeated
88 names; the registry lives on the root archive so that nested
89 archives share a single namespace. Each backend chooses how to
90 encode ``version > 1`` on disk: FITS uses the FITS ``EXTVER``
91 keyword without modifying the extension name, while hierarchical
92 backends can append ``_{version}`` to the leaf component of the path.
93 """
94 version = self._name_versions.get(name, 0) + 1
95 self._name_versions[name] = version
96 return name, version
98 @abstractmethod
99 def serialize_direct[T: pydantic.BaseModel | None](
100 self, name: str, serializer: Callable[[OutputArchive], T]
101 ) -> T:
102 """Use a serializer function to save a nested object.
104 Parameters
105 ----------
106 name
107 Attribute of the paired Pydantic model that will be assigned the
108 result of this call. If it will not be assigned to a direct
109 attribute, it may be a JSON Pointer path (relative to the paired
110 Pydantic model) to the location where it will be added.
111 serializer
112 Callable that takes an `~lsst.serialization.OutputArchive` and
113 returns a Pydantic model. This will be passed a new
114 `~lsst.serialization.OutputArchive` that automatically prepends
115 ``{name}/`` (and any root path added by this archive) to names
116 passed to it, so the ``serializer`` does not need to know where it
117 appears in the overall tree.
119 Returns
120 -------
121 T
122 Result of the call to the serializer.
123 """
124 raise NotImplementedError()
126 @abstractmethod
127 def serialize_pointer[T: ArchiveTree](
128 self, name: str, serializer: Callable[[OutputArchive], T], key: Hashable
129 ) -> T | P:
130 """Use a serializer function to save a nested object that may be
131 referenced in multiple locations in the same archive.
133 Parameters
134 ----------
135 name
136 Attribute of the paired Pydantic model that will be assigned the
137 result of this call. If it will not be assigned to a direct
138 attribute, it may be a JSON Pointer path (relative to the paired
139 Pydantic model) to the location where it will be added.
140 serializer
141 Callable that takes an `~lsst.serialization.OutputArchive` and
142 returns a Pydantic model. This will be passed a new
143 `~lsst.serialization.OutputArchive` that automatically prepends
144 ``{name}/`` (and any root path added by this archive) to names
145 passed to it, so the ``serializer`` does not need to know where it
146 appears in the overall tree.
147 key
148 A unique identifier for the in-memory object the serializer saves,
149 e.g. a call to the built-in `id` function.
151 Returns
152 -------
153 T | P
154 Either the result of the call to the serializer, or a Pydantic
155 model that can be considered a reference to it and added to a
156 larger model in its place.
157 """
158 # Since Pydantic doesn't provide us a good way to "dereference" a JSON
159 # Pointer (i.e. traversing the tree to extract the original model), it
160 # is probably easier to implement an `InputArchive` for the case where
161 # the `~lsst.serialization.OutputArchive` opts to stuff all pointer
162 # serializations into a standard location outside the user-controlled
163 # Pydantic model tree, and always returned a JSON pointer to that
164 # standard location from this function.
165 raise NotImplementedError()
167 @abstractmethod
168 def serialize_frame_set[T: ArchiveTree](
169 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
170 ) -> T | P:
171 """Serialize a frame set and make it available to objects saved later.
173 Parameters
174 ----------
175 name
176 Attribute of the paired Pydantic model that will be assigned the
177 result of this call. If it will not be assigned to a direct
178 attribute, it may be a JSON Pointer path (relative to the paired
179 Pydantic model) to the location where it will be added.
180 frame_set
181 The frame set being saved. This will be returned in later calls
182 to `iter_frame_sets`, along with the returned reference object.
183 serializer
184 Callable that takes an `~lsst.serialization.OutputArchive` and
185 returns a Pydantic model. This will be passed a new
186 `~lsst.serialization.OutputArchive` that automatically prepends
187 ``{name}/`` (and any root path added by this archive) to names
188 passed to it, so the ``serializer`` does not need to know where it
189 appears in the overall tree.
190 key
191 A unique identifier for the in-memory object the serializer saves,
192 e.g. a call to the built-in `id` function.
194 Returns
195 -------
196 T | P
197 Either the result of the call to the serializer, or a Pydantic
198 model that can be considered a reference to it and added to a
199 larger model in its place.
200 """
201 raise NotImplementedError()
203 @abstractmethod
204 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, P]]:
205 """Iterate over the frame sets already serialized to this archive.
207 Yields
208 ------
209 frame_set
210 A frame set that has already been written to this archive.
211 reference
212 An implementation-specific reference model that points to the
213 frame set.
214 """
215 raise NotImplementedError()
217 @abstractmethod
218 def add_array(
219 self,
220 array: np.ndarray,
221 *,
222 name: str | None = None,
223 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
224 tile_shape: tuple[int, ...] | None = None,
225 options_name: str | None = None,
226 ) -> ArrayReferenceModel | InlineArrayModel:
227 """Add an array to the archive.
229 Parameters
230 ----------
231 array
232 Array to save.
233 name
234 Name of the array. This should generally be the name of the
235 Pydantic model attribute to which the result will be assigned. It
236 may be left `None` if there is only one [structured] array or
237 table in a nested object that is being saved.
238 update_header
239 A callback that will be given the FITS header for the HDU
240 containing this array in order to add keys to it. This callback
241 may be provided but will not be called if the output format is not
242 FITS.
243 tile_shape
244 The recommended shape of each tile if the implementation will save
245 the array in distinct tiles for faster subarray retrieval.
246 This is a hint; implementations are not required to use this value.
247 options_name
248 Use the options (e.g. for compression) associated with this name
249 when saving this array.
251 Returns
252 -------
253 `~lsst.images.serialization.ArrayReferenceModel` |\
254 `~lsst.images.serialization.InlineArrayModel`
255 A Pydantic model that references or holds the stored array.
256 """
257 raise NotImplementedError()
259 @abstractmethod
260 def add_table(
261 self,
262 table: astropy.table.Table,
263 *,
264 name: str | None = None,
265 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
266 ) -> TableModel:
267 """Add a table to the archive.
269 Parameters
270 ----------
271 table
272 Table to save.
273 name
274 Name of the table. This should generally be the name of the
275 Pydantic model attribute to which the result will be assigned. It
276 may be left `None` if there is only one [structured] array or
277 table in a nested object that is being saved.
278 update_header
279 A callback that will be given the FITS header for the HDU
280 containing this table in order to add keys to it. This callback
281 may be provided but will not be called if the output format is not
282 FITS.
284 Returns
285 -------
286 TableModel
287 A Pydantic model that represents the table.
288 """
289 raise NotImplementedError()
291 @abstractmethod
292 def add_structured_array(
293 self,
294 array: np.ndarray,
295 *,
296 name: str | None = None,
297 units: Mapping[str, astropy.units.Unit] | None = None,
298 descriptions: Mapping[str, str] | None = None,
299 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
300 ) -> TableModel:
301 """Add a table to the archive.
303 Parameters
304 ----------
305 name
306 Attribute of the paired Pydantic model that will be assigned the
307 result of this call. If it will not be assigned to a direct
308 attribute, it may be a JSON Pointer path (relative to the paired
309 Pydantic model) to the location where it will be added.
310 array
311 A structured numpy array.
312 name
313 Name of the array. This should generally be the name of the
314 Pydantic model attribute to which the result will be assigned. It
315 may be left `None` if there is only one [structured] array or
316 table in a nested object that is being saved.
317 units
318 A mapping of units for columns. Need not be complete.
319 descriptions
320 A mapping of descriptions for columns. Need not be complete.
321 update_header
322 A callback that will be given the FITS header for the HDU
323 containing this table in order to add keys to it. This callback
324 may be provided but will not be called if the output format is not
325 FITS.
327 Returns
328 -------
329 TableModel
330 A Pydantic model that represents the table.
331 """
332 raise NotImplementedError()
335class NestedOutputArchive[P: pydantic.BaseModel](OutputArchive[P]):
336 """A proxy output archive that joins a root path into all names before
337 delegating back to its parent archive.
339 This is intended to be used in the implementation of most
340 `~lsst.serialization.OutputArchive.serialize_direct` and
341 `~lsst.serialization.OutputArchive.serialize_pointer` implementations.
343 Parameters
344 ----------
345 root
346 Root of all JSON Pointer paths. Should include a leading slash (as we
347 always use absolute JSON Pointers) but no trailing slash.
348 parent
349 Parent output archive to delegate to.
350 """
352 def __init__(self, root: str, parent: OutputArchive):
353 super().__init__()
354 self._root = root
355 self._parent = parent
357 def serialize_direct[T: pydantic.BaseModel | None](
358 self, name: str, serializer: Callable[[OutputArchive[P]], T]
359 ) -> T:
360 return self._parent.serialize_direct(self._join_path(name), serializer)
362 def serialize_pointer[T: ArchiveTree](
363 self, name: str, serializer: Callable[[OutputArchive[P]], T], key: Hashable
364 ) -> T | P:
365 return self._parent.serialize_pointer(self._join_path(name), serializer, key)
367 def serialize_frame_set[T: ArchiveTree](
368 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
369 ) -> T | P:
370 return self._parent.serialize_frame_set(self._join_path(name), frame_set, serializer, key)
372 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, P]]:
373 return self._parent.iter_frame_sets()
375 def add_array(
376 self,
377 array: np.ndarray,
378 *,
379 name: str | None = None,
380 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
381 tile_shape: tuple[int, ...] | None = None,
382 options_name: str | None = None,
383 ) -> ArrayReferenceModel | InlineArrayModel:
384 return self._parent.add_array(
385 array,
386 name=self._join_path(name),
387 update_header=update_header,
388 tile_shape=tile_shape,
389 options_name=options_name,
390 )
392 def add_table(
393 self,
394 table: astropy.table.Table,
395 *,
396 name: str | None = None,
397 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
398 ) -> TableModel:
399 return self._parent.add_table(table, name=self._join_path(name), update_header=update_header)
401 def add_structured_array(
402 self,
403 array: np.ndarray,
404 *,
405 name: str | None = None,
406 units: Mapping[str, astropy.units.Unit] | None = None,
407 descriptions: Mapping[str, str] | None = None,
408 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
409 ) -> TableModel:
410 return self._parent.add_structured_array(
411 array,
412 name=self._join_path(name),
413 units=units,
414 descriptions=descriptions,
415 update_header=update_header,
416 )
418 def _join_path(self, name: str | None) -> str:
419 return f"{self._root}/{name}" if name is not None else self._root