Coverage for python/lsst/images/serialization/_output_archive.py: 91%
57 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-08 02:03 -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 array
306 A structured numpy array.
307 name
308 Name of the array. This should generally be the name of the
309 Pydantic model attribute to which the result will be assigned. It
310 may be left `None` if there is only one [structured] array or
311 table in a nested object that is being saved.
312 units
313 A mapping of units for columns. Need not be complete.
314 descriptions
315 A mapping of descriptions for columns. Need not be complete.
316 update_header
317 A callback that will be given the FITS header for the HDU
318 containing this table in order to add keys to it. This callback
319 may be provided but will not be called if the output format is not
320 FITS.
322 Returns
323 -------
324 TableModel
325 A Pydantic model that represents the table.
326 """
327 raise NotImplementedError()
330class NestedOutputArchive[P: pydantic.BaseModel](OutputArchive[P]):
331 """A proxy output archive that joins a root path into all names before
332 delegating back to its parent archive.
334 This is intended to be used in the implementation of most
335 `~lsst.serialization.OutputArchive.serialize_direct` and
336 `~lsst.serialization.OutputArchive.serialize_pointer` implementations.
338 Parameters
339 ----------
340 root
341 Root of all JSON Pointer paths. Should include a leading slash (as we
342 always use absolute JSON Pointers) but no trailing slash.
343 parent
344 Parent output archive to delegate to.
345 """
347 def __init__(self, root: str, parent: OutputArchive) -> None:
348 super().__init__()
349 self._root = root
350 self._parent = parent
352 def serialize_direct[T: pydantic.BaseModel | None](
353 self, name: str, serializer: Callable[[OutputArchive[P]], T]
354 ) -> T:
355 return self._parent.serialize_direct(self._join_path(name), serializer)
357 def serialize_pointer[T: ArchiveTree](
358 self, name: str, serializer: Callable[[OutputArchive[P]], T], key: Hashable
359 ) -> T | P:
360 return self._parent.serialize_pointer(self._join_path(name), serializer, key)
362 def serialize_frame_set[T: ArchiveTree](
363 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
364 ) -> T | P:
365 return self._parent.serialize_frame_set(self._join_path(name), frame_set, serializer, key)
367 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, P]]:
368 return self._parent.iter_frame_sets()
370 def add_array(
371 self,
372 array: np.ndarray,
373 *,
374 name: str | None = None,
375 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
376 tile_shape: tuple[int, ...] | None = None,
377 options_name: str | None = None,
378 ) -> ArrayReferenceModel | InlineArrayModel:
379 return self._parent.add_array(
380 array,
381 name=self._join_path(name),
382 update_header=update_header,
383 tile_shape=tile_shape,
384 options_name=options_name,
385 )
387 def add_table(
388 self,
389 table: astropy.table.Table,
390 *,
391 name: str | None = None,
392 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
393 ) -> TableModel:
394 return self._parent.add_table(table, name=self._join_path(name), update_header=update_header)
396 def add_structured_array(
397 self,
398 array: np.ndarray,
399 *,
400 name: str | None = None,
401 units: Mapping[str, astropy.units.Unit] | None = None,
402 descriptions: Mapping[str, str] | None = None,
403 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
404 ) -> TableModel:
405 return self._parent.add_structured_array(
406 array,
407 name=self._join_path(name),
408 units=units,
409 descriptions=descriptions,
410 update_header=update_header,
411 )
413 def _join_path(self, name: str | None) -> str:
414 return f"{self._root}/{name}" if name is not None else self._root