Coverage for python/lsst/images/serialization/_output_archive.py: 33%
66 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 07:45 +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 "NestedOutputArchive",
16 "OutputArchive",
17)
19from abc import ABC, abstractmethod
20from collections.abc import Callable, Hashable, Iterator, Mapping
21from typing import TYPE_CHECKING, Any, 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 (
31 ArchiveTree,
32 ButlerInfo,
33 MetadataValue,
34 no_header_updates,
35 warn_for_development_schemas,
36)
37from ._tables import TableModel
39if TYPE_CHECKING:
40 from .._transforms import FrameSet
42# This pre-python-3.12 declaration is needed by Sphinx (probably the
43# autodoc-typehints plugin.
44P = TypeVar("P", bound=pydantic.BaseModel)
47class OutputArchive[P](ABC):
48 """Abstract interface for writing to a file format.
50 Notes
51 -----
52 An output archive instance is assumed to be paired with a Pydantic model
53 that represents a JSON tree, with the archive used to serialize data that
54 is not natively JSON into data that is (which may just be a reference to
55 binary data stored elsewhere in the file). The archive doesn't actually
56 hold that model instance because we don't want to assume it can be built
57 via default-initialization and assignment, and because we'd prefer to avoid
58 making the output archive generic over the model type. It is expected that
59 most concrete archive implementations will accept the paired model in some
60 sort of finalization method in order to write it into the file, but this is
61 not part of the base class interface.
62 """
64 def __init__(self) -> None:
65 self._name_versions: dict[str, int] = {}
66 """Per-name occurrence count, used by `_register_name` to disambiguate
67 repeated logical names within a single write (e.g. each operand of a
68 `SumField` calling ``add_array(name="data")`` from the same nested
69 archive).
70 """
72 def serialize_root(
73 self,
74 obj: Any,
75 metadata: dict[str, MetadataValue] | None = None,
76 butler_info: ButlerInfo | None = None,
77 ) -> ArchiveTree:
78 """Serialize ``obj`` to a root tree, apply write-time overrides, and
79 warn if the tree uses a development schema.
81 Every backend's top-level write funnel calls this so the metadata and
82 butler overrides and the development-schema warning are applied
83 uniformly, regardless of file format.
85 Parameters
86 ----------
87 obj
88 Object with a ``serialize`` method (and an optional
89 ``_archive_default_name``) to serialize.
90 metadata
91 Extra metadata to merge into the tree, or `None`.
92 butler_info
93 Butler information to set on the tree, or `None`.
95 Returns
96 -------
97 `~lsst.images.serialization.ArchiveTree`
98 The serialized root tree.
99 """
100 name = getattr(obj, "_archive_default_name", None)
101 tree = self.serialize_direct(name, obj.serialize) if name is not None else obj.serialize(self)
102 if metadata is not None:
103 tree.metadata.update(metadata)
104 if butler_info is not None:
105 tree.butler_info = butler_info
106 warn_for_development_schemas(tree)
107 return tree
109 def _register_name(self, name: str) -> tuple[str, int]:
110 """Return the input name and its 1-based occurrence count.
112 Parameters
113 ----------
114 name
115 The logical archive name being saved (typically the absolute
116 archive path of an array, table, or pointer target).
118 Returns
119 -------
120 name : `str`
121 The input name, returned unchanged so that the caller controls
122 how the version is rendered into the on-disk identifier.
123 version : `int`
124 ``1`` the first time a given name is registered, then ``2``,
125 ``3`` and so on for subsequent calls with the same name.
127 Notes
128 -----
129 Backends should call this from `add_array`, `add_table`,
130 `add_structured_array`, and `serialize_pointer` to detect repeated
131 names; the registry lives on the root archive so that nested
132 archives share a single namespace. Each backend chooses how to
133 encode ``version > 1`` on disk: FITS uses the FITS ``EXTVER``
134 keyword without modifying the extension name, while hierarchical
135 backends can append ``_{version}`` to the leaf component of the path.
136 """
137 version = self._name_versions.get(name, 0) + 1
138 self._name_versions[name] = version
139 return name, version
141 @abstractmethod
142 def serialize_direct[T: pydantic.BaseModel | None](
143 self, name: str, serializer: Callable[[OutputArchive], T]
144 ) -> T:
145 """Use a serializer function to save a nested object.
147 Parameters
148 ----------
149 name
150 Attribute of the paired Pydantic model that will be assigned the
151 result of this call. If it will not be assigned to a direct
152 attribute, it may be a JSON Pointer path (relative to the paired
153 Pydantic model) to the location where it will be added.
154 serializer
155 Callable that takes an `~lsst.serialization.OutputArchive` and
156 returns a Pydantic model. This will be passed a new
157 `~lsst.serialization.OutputArchive` that automatically prepends
158 ``{name}/`` (and any root path added by this archive) to names
159 passed to it, so the ``serializer`` does not need to know where it
160 appears in the overall tree.
162 Returns
163 -------
164 T
165 Result of the call to the serializer.
166 """
167 raise NotImplementedError()
169 @abstractmethod
170 def serialize_pointer[T: ArchiveTree](
171 self, name: str, serializer: Callable[[OutputArchive], T], key: Hashable
172 ) -> T | P:
173 """Use a serializer function to save a nested object that may be
174 referenced in multiple locations in the same archive.
176 Parameters
177 ----------
178 name
179 Attribute of the paired Pydantic model that will be assigned the
180 result of this call. If it will not be assigned to a direct
181 attribute, it may be a JSON Pointer path (relative to the paired
182 Pydantic model) to the location where it will be added.
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 # Since Pydantic doesn't provide us a good way to "dereference" a JSON
202 # Pointer (i.e. traversing the tree to extract the original model), it
203 # is probably easier to implement an `InputArchive` for the case where
204 # the `~lsst.serialization.OutputArchive` opts to stuff all pointer
205 # serializations into a standard location outside the user-controlled
206 # Pydantic model tree, and always returned a JSON pointer to that
207 # standard location from this function.
208 raise NotImplementedError()
210 @abstractmethod
211 def serialize_frame_set[T: ArchiveTree](
212 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
213 ) -> T | P:
214 """Serialize a frame set and make it available to objects saved later.
216 Parameters
217 ----------
218 name
219 Attribute of the paired Pydantic model that will be assigned the
220 result of this call. If it will not be assigned to a direct
221 attribute, it may be a JSON Pointer path (relative to the paired
222 Pydantic model) to the location where it will be added.
223 frame_set
224 The frame set being saved. This will be returned in later calls
225 to `iter_frame_sets`, along with the returned reference object.
226 serializer
227 Callable that takes an `~lsst.serialization.OutputArchive` and
228 returns a Pydantic model. This will be passed a new
229 `~lsst.serialization.OutputArchive` that automatically prepends
230 ``{name}/`` (and any root path added by this archive) to names
231 passed to it, so the ``serializer`` does not need to know where it
232 appears in the overall tree.
233 key
234 A unique identifier for the in-memory object the serializer saves,
235 e.g. a call to the built-in `id` function.
237 Returns
238 -------
239 T | P
240 Either the result of the call to the serializer, or a Pydantic
241 model that can be considered a reference to it and added to a
242 larger model in its place.
243 """
244 raise NotImplementedError()
246 @abstractmethod
247 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, P]]:
248 """Iterate over the frame sets already serialized to this archive.
250 Yields
251 ------
252 frame_set
253 A frame set that has already been written to this archive.
254 reference
255 An implementation-specific reference model that points to the
256 frame set.
257 """
258 raise NotImplementedError()
260 @abstractmethod
261 def add_array(
262 self,
263 array: np.ndarray,
264 *,
265 name: str | None = None,
266 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
267 tile_shape: tuple[int, ...] | None = None,
268 options_name: str | None = None,
269 ) -> ArrayReferenceModel | InlineArrayModel:
270 """Add an array to the archive.
272 Parameters
273 ----------
274 array
275 Array to save.
276 name
277 Name of the array. This should generally be the name of the
278 Pydantic model attribute to which the result will be assigned. It
279 may be left `None` if there is only one [structured] array or
280 table in a nested object that is being saved.
281 update_header
282 A callback that will be given the FITS header for the HDU
283 containing this array in order to add keys to it. This callback
284 may be provided but will not be called if the output format is not
285 FITS.
286 tile_shape
287 The recommended shape of each tile if the implementation will save
288 the array in distinct tiles for faster subarray retrieval.
289 This is a hint; implementations are not required to use this value.
290 options_name
291 Use the options (e.g. for compression) associated with this name
292 when saving this array.
294 Returns
295 -------
296 `~lsst.images.serialization.ArrayReferenceModel` |\
297 `~lsst.images.serialization.InlineArrayModel`
298 A Pydantic model that references or holds the stored array.
299 """
300 raise NotImplementedError()
302 @abstractmethod
303 def add_table(
304 self,
305 table: astropy.table.Table,
306 *,
307 name: str | None = None,
308 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
309 ) -> TableModel:
310 """Add a table to the archive.
312 Parameters
313 ----------
314 table
315 Table to save.
316 name
317 Name of the table. This should generally be the name of the
318 Pydantic model attribute to which the result will be assigned. It
319 may be left `None` if there is only one [structured] array or
320 table in a nested object that is being saved.
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()
334 @abstractmethod
335 def add_structured_array(
336 self,
337 array: np.ndarray,
338 *,
339 name: str | None = None,
340 units: Mapping[str, astropy.units.Unit] | None = None,
341 descriptions: Mapping[str, str] | None = None,
342 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
343 ) -> TableModel:
344 """Add a table to the archive.
346 Parameters
347 ----------
348 array
349 A structured numpy array.
350 name
351 Name of the array. This should generally be the name of the
352 Pydantic model attribute to which the result will be assigned. It
353 may be left `None` if there is only one [structured] array or
354 table in a nested object that is being saved.
355 units
356 A mapping of units for columns. Need not be complete.
357 descriptions
358 A mapping of descriptions for columns. Need not be complete.
359 update_header
360 A callback that will be given the FITS header for the HDU
361 containing this table in order to add keys to it. This callback
362 may be provided but will not be called if the output format is not
363 FITS.
365 Returns
366 -------
367 TableModel
368 A Pydantic model that represents the table.
369 """
370 raise NotImplementedError()
373class NestedOutputArchive[P: pydantic.BaseModel](OutputArchive[P]):
374 """A proxy output archive that joins a root path into all names before
375 delegating back to its parent archive.
377 This is intended to be used in the implementation of most
378 `~lsst.serialization.OutputArchive.serialize_direct` and
379 `~lsst.serialization.OutputArchive.serialize_pointer` implementations.
381 Parameters
382 ----------
383 root
384 Root of all JSON Pointer paths. Should include a leading slash (as we
385 always use absolute JSON Pointers) but no trailing slash.
386 parent
387 Parent output archive to delegate to.
388 """
390 def __init__(self, root: str, parent: OutputArchive) -> None:
391 super().__init__()
392 self._root = root
393 self._parent = parent
395 def serialize_direct[T: pydantic.BaseModel | None](
396 self, name: str, serializer: Callable[[OutputArchive[P]], T]
397 ) -> T:
398 return self._parent.serialize_direct(self._join_path(name), serializer)
400 def serialize_pointer[T: ArchiveTree](
401 self, name: str, serializer: Callable[[OutputArchive[P]], T], key: Hashable
402 ) -> T | P:
403 return self._parent.serialize_pointer(self._join_path(name), serializer, key)
405 def serialize_frame_set[T: ArchiveTree](
406 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
407 ) -> T | P:
408 return self._parent.serialize_frame_set(self._join_path(name), frame_set, serializer, key)
410 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, P]]:
411 return self._parent.iter_frame_sets()
413 def add_array(
414 self,
415 array: np.ndarray,
416 *,
417 name: str | None = None,
418 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
419 tile_shape: tuple[int, ...] | None = None,
420 options_name: str | None = None,
421 ) -> ArrayReferenceModel | InlineArrayModel:
422 return self._parent.add_array(
423 array,
424 name=self._join_path(name),
425 update_header=update_header,
426 tile_shape=tile_shape,
427 options_name=options_name,
428 )
430 def add_table(
431 self,
432 table: astropy.table.Table,
433 *,
434 name: str | None = None,
435 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
436 ) -> TableModel:
437 return self._parent.add_table(table, name=self._join_path(name), update_header=update_header)
439 def add_structured_array(
440 self,
441 array: np.ndarray,
442 *,
443 name: str | None = None,
444 units: Mapping[str, astropy.units.Unit] | None = None,
445 descriptions: Mapping[str, str] | None = None,
446 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
447 ) -> TableModel:
448 return self._parent.add_structured_array(
449 array,
450 name=self._join_path(name),
451 units=units,
452 descriptions=descriptions,
453 update_header=update_header,
454 )
456 def _join_path(self, name: str | None) -> str:
457 return f"{self._root}/{name}" if name is not None else self._root