Coverage for python / lsst / images / fits / _output_archive.py: 21%
205 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-31 02:05 -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__ = ("FitsOutputArchive", "write")
16import dataclasses
17from collections import Counter
18from collections.abc import Callable, Hashable, Iterator, Mapping
19from contextlib import contextmanager
20from typing import Any, Self
22import astropy.io.fits
23import astropy.table
24import numpy as np
25import pydantic
27from .._transforms import FrameSet
28from ..serialization import (
29 ArchiveTree,
30 ArrayReferenceModel,
31 ButlerInfo,
32 ColumnDefinitionModel,
33 MetadataValue,
34 NestedOutputArchive,
35 NumberType,
36 OutputArchive,
37 TableCellReferenceModel,
38 TableModel,
39 no_header_updates,
40)
41from ._common import ExtensionHDU, ExtensionKey, FitsCompressionOptions, FitsOpaqueMetadata
44def write(
45 obj: Any,
46 filename: str,
47 compression_options: Mapping[str, FitsCompressionOptions | None] | None = None,
48 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
49 compression_seed: int | None = None,
50 metadata: dict[str, MetadataValue] | None = None,
51 butler_info: ButlerInfo | None = None,
52) -> Any:
53 """Write an object with a ``serialize`` method to a FITS file.
55 Parameters
56 ----------
57 filename
58 Name of the file to write to. Must not already exist.
59 compression_options
60 Options for how to compress the FITS file, keyed by the name of
61 the attribute (with JSON pointer ``/`` separators for nested
62 attributes).
63 update_header
64 A callback that will be given the primary HDU FITS header and an
65 opportunity to modify it.
66 compression_seed
67 A FITS tile compression seed to use whenever the configured
68 compression seed is `None` or (for backwards compatibility) ``0``.
69 This value is then incremented every time it is used.
70 metadata
71 Additional metadata to save with the object. This will override any
72 flexible metadata carried by the object itself with the same keys.
73 butler_info
74 Butler information to store in the file.
76 Returns
77 -------
78 `.serialization.ArchiveTree`
79 The serialized representation of the object.
80 """
81 opaque_metadata = getattr(obj, "_opaque_metadata", None)
82 name = getattr(obj, "_archive_default_name", None)
83 with FitsOutputArchive.open(
84 filename,
85 compression_options=compression_options,
86 opaque_metadata=opaque_metadata,
87 update_header=update_header,
88 compression_seed=compression_seed,
89 ) as archive:
90 tree = archive.serialize_direct(name, obj.serialize) if name is not None else obj.serialize(archive)
91 if metadata is not None:
92 tree.metadata.update(metadata)
93 if butler_info is not None:
94 tree.butler_info = butler_info
95 archive.add_tree(tree)
96 return tree
99class FitsOutputArchive(OutputArchive[TableCellReferenceModel]):
100 """An implementation of the `.serialization.OutputArchive` interface that
101 writes to FITS files.
103 Instances of this class should only be constructed via the `open`
104 context manager.
105 """
107 def __init__(
108 self,
109 hdu_list: astropy.io.fits.HDUList,
110 compression_options: Mapping[str, FitsCompressionOptions | None] | None = None,
111 opaque_metadata: Any = None,
112 compression_seed: int | None = None,
113 ):
114 # JSON blobs for objects we've saved as pointers:
115 self._pointer_targets: list[bytes] = []
116 # Mapping from user provided key (e.g. id(some object)) to a table
117 # pointer to where we actually saved it:
118 self._pointers_by_key: dict[Hashable, TableCellReferenceModel] = {}
119 self._hdu_list = hdu_list
120 self._primary_hdu = astropy.io.fits.PrimaryHDU()
121 # TODO: add subformat description and version to primary HDU.
122 self._primary_hdu.header.set("INDXADDR", 0, "Offset in bytes to the HDU index.")
123 self._primary_hdu.header.set("INDXSIZE", 0, "Size of the HDU index.")
124 self._primary_hdu.header.set("JSONADDR", 0, "Offset in bytes to the JSON tree HDU.")
125 self._primary_hdu.header.set("JSONSIZE", 0, "Size of the JSON tree HDU.")
126 self._hdus_by_name = Counter[str]()
127 self._compression_options = dict(compression_options) if compression_options is not None else {}
128 self._compression_seed = compression_seed
129 self._opaque_metadata = (
130 opaque_metadata if isinstance(opaque_metadata, FitsOpaqueMetadata) else FitsOpaqueMetadata()
131 )
132 if (opaque_primary_header := self._opaque_metadata.headers.get(ExtensionKey())) is not None:
133 self._primary_hdu.header.extend(opaque_primary_header)
134 self._hdu_list.append(self._primary_hdu)
135 self._json_hdu_added: bool = False
136 self._frame_sets: list[tuple[FrameSet, TableCellReferenceModel]] = []
138 @classmethod
139 @contextmanager
140 def open(
141 cls,
142 filename: str,
143 compression_options: Mapping[str, FitsCompressionOptions | None] | None = None,
144 opaque_metadata: Any = None,
145 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
146 compression_seed: int | None = None,
147 ) -> Iterator[Self]:
148 """Create an output archive that writes to the given file.
150 Parameters
151 ----------
152 filename
153 Name of the file to write to. Must not already exist.
154 compression_options
155 Options for how to compress the FITS file, keyed by the name of
156 the attribute (with JSON pointer ``/`` separators for nested
157 attributes).
158 opaque_metadata
159 Metadata read from an input archive along with the object being
160 written now. Ignored if the metadata is not from a FITS archive.
161 update_header
162 A callback that will be given the primary HDU FITS header and an
163 opportunity to modify it.
164 compression_seed
165 A FITS tile compression seed to use whenever the configured
166 compression seed is `None` or (for backwards compatibility) ``0``.
167 This value is then incremented every time it is used.
169 Returns
170 -------
171 `contextlib.AbstractContextManager` [`FitsOutputArchive`]
172 A context manager that returns a `FitsOutputArchive` when entered.
173 """
174 with astropy.io.fits.open(filename, mode="append") as hdu_list:
175 if hdu_list:
176 raise OSError(f"File {filename!r} already exists.")
177 archive = cls(hdu_list, compression_options, opaque_metadata, compression_seed=compression_seed)
178 update_header(hdu_list[0].header)
179 yield archive
180 if not archive._json_hdu_added:
181 raise RuntimeError("Write context exited without 'add_tree' being called.")
182 hdu_list.flush()
183 # This multi-open dance is necessary to get Astropy to tell us the
184 # byte addresses of the HDUs. Hopefully we can get an upstream change
185 # make this unnecessary at some point.
186 with astropy.io.fits.open(filename, mode="readonly", disable_image_compression=True) as hdu_list:
187 index_hdu = cls._make_index_table(hdu_list)
188 with astropy.io.fits.open(filename, mode="append") as hdu_list:
189 hdu_list.append(index_hdu)
190 json_bytes = _HDUBytes.from_index_row(index_hdu.data[-1])
191 index_bytes = _HDUBytes.from_write_hdu(index_hdu)
192 # Update the primary HDU with the address and size of the index and
193 # JSON HDUs, and rewrite just that. We do this write manually, since
194 # astropy's docs on its 'update' mode are scarce and it's not obvious
195 # whether we can guarantee it won't rewrite the whole file if we edit
196 # the primary header.
197 archive._primary_hdu.header["INDXADDR"] = index_bytes.header_address
198 archive._primary_hdu.header["INDXSIZE"] = index_bytes.size
199 archive._primary_hdu.header["JSONADDR"] = json_bytes.header_address
200 archive._primary_hdu.header["JSONSIZE"] = json_bytes.size
201 with open(filename, "r+b") as stream:
202 stream.write(archive._primary_hdu.header.tostring().encode())
204 def serialize_direct[T: pydantic.BaseModel](
205 self, name: str, serializer: Callable[[OutputArchive[TableCellReferenceModel]], T]
206 ) -> T:
207 nested = NestedOutputArchive[TableCellReferenceModel](name, self)
208 return serializer(nested)
210 def serialize_pointer[T: ArchiveTree](
211 self, name: str, serializer: Callable[[OutputArchive[TableCellReferenceModel]], T], key: Hashable
212 ) -> TableCellReferenceModel:
213 if (pointer := self._pointers_by_key.get(key)) is not None:
214 return pointer
215 pointer = TableCellReferenceModel(
216 source="fits:JSON", column="JSON", row=len(self._pointer_targets) + 1
217 )
218 model = self.serialize_direct("", serializer)
219 self._pointer_targets.append(model.model_dump_json().encode())
220 self._pointers_by_key[key] = pointer
221 return pointer
223 def serialize_frame_set[T: ArchiveTree](
224 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
225 ) -> TableCellReferenceModel:
226 # Docstring inherited.
227 pointer = self.serialize_pointer(name, serializer, key)
228 self._frame_sets.append((frame_set, pointer))
229 return pointer
231 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, TableCellReferenceModel]]:
232 return iter(self._frame_sets)
234 def add_array(
235 self,
236 array: np.ndarray,
237 *,
238 name: str | None = None,
239 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
240 ) -> ArrayReferenceModel:
241 if name is None:
242 raise RuntimeError("Cannot save array with name=None unless it is nested.")
243 extname = name.upper()
244 hdu = self._opaque_metadata.maybe_use_precompressed(extname)
245 if hdu is None:
246 if (compression_options := self._get_compression_options(name)) is not None:
247 hdu = compression_options.make_hdu(array, name=extname)
248 else:
249 hdu = astropy.io.fits.ImageHDU(array, name=extname)
250 key = self._add_hdu(hdu, update_header)
251 return ArrayReferenceModel(
252 source=str(key),
253 shape=list(array.shape),
254 datatype=NumberType.from_numpy(array.dtype),
255 )
257 def add_table(
258 self,
259 table: astropy.table.Table,
260 *,
261 name: str | None = None,
262 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
263 ) -> TableModel:
264 if name is None:
265 raise RuntimeError("Cannot save table with name=None unless it is nested.")
266 extname = name.upper()
267 hdu: astropy.io.fits.BinTableHDU = astropy.io.fits.table_to_hdu(table, name=extname)
268 columns = ColumnDefinitionModel.from_record_dtype(hdu.data.dtype)
269 for c in columns:
270 c.update_from_table(table)
271 key = self._add_hdu(hdu, update_header)
272 return TableModel(source=str(key), columns=columns)
274 def add_structured_array(
275 self,
276 array: np.ndarray,
277 *,
278 name: str | None = None,
279 units: Mapping[str, astropy.units.Unit] | None = None,
280 descriptions: Mapping[str, str] | None = None,
281 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
282 ) -> TableModel:
283 if name is None:
284 raise RuntimeError("Cannot save structured array with name=None unless it is nested.")
285 extname = name.upper()
286 hdu = astropy.io.fits.BinTableHDU(array, name=extname)
287 columns = ColumnDefinitionModel.from_record_dtype(hdu.data.dtype)
288 if units is not None:
289 for c in columns:
290 c.unit = units.get(c.name)
291 if descriptions is not None:
292 for c in columns:
293 c.description = descriptions.get(c.name, "")
294 key = self._add_hdu(hdu, update_header)
295 return TableModel(source=str(key), columns=columns)
297 def _add_hdu(
298 self,
299 hdu: astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU | astropy.io.fits.BinTableHDU,
300 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
301 ) -> ExtensionKey:
302 n_hdus = self._hdus_by_name.get(hdu.name, 0)
303 key = ExtensionKey(hdu.name, n_hdus + 1)
304 if n_hdus:
305 hdu.header["EXTVER"] = key.ver
306 self._hdus_by_name[hdu.name] += 1
307 update_header(hdu.header)
308 if (opaque_headers := self._opaque_metadata.headers.get(key)) is not None:
309 hdu.header.extend(opaque_headers)
310 self._hdu_list.append(hdu)
311 return key
313 def add_tree(self, tree: ArchiveTree) -> None:
314 """Write the JSON tree to the archive.
316 This method must be called exactly once, just before the `open` context
317 is exited.
319 Parameters
320 ----------
321 tree
322 Pydantic model that represents the tree.
323 """
324 json_hdu = astropy.io.fits.BinTableHDU.from_columns(
325 [astropy.io.fits.Column("JSON", "PB")],
326 nrows=len(self._pointer_targets) + 1,
327 name="JSON",
328 )
329 json_hdu.data[0]["JSON"] = np.frombuffer(tree.model_dump_json().encode(), dtype=np.byte)
330 for n, json_target_data in enumerate(self._pointer_targets):
331 json_hdu.data[n + 1]["JSON"] = np.frombuffer(json_target_data, dtype=np.byte)
332 self._hdu_list.append(json_hdu)
333 self._json_hdu_added = True
335 def _get_compression_options(self, name: str) -> FitsCompressionOptions | None:
336 result = self._compression_options.get(name, FitsCompressionOptions.DEFAULT)
337 if result is None or result.quantization is None:
338 return result
339 if self._compression_seed is not None and not result.quantization.seed:
340 result = result.model_copy(
341 update={
342 "quantization": result.quantization.model_copy(update={"seed": self._compression_seed})
343 }
344 )
345 self._compression_seed += 1
346 if self._compression_seed > 10000:
347 self._compression_seed = 1
348 # MyPy can tell that result.quantization is not None in the 'if', but
349 # forgets that by this 'else':
350 elif result.quantization.seed is None: # type: ignore[union-attr]
351 raise RuntimeError("No quantization seed provided.")
352 return result
354 @staticmethod
355 def _make_index_table(hdu_list: astropy.io.fits.HDUList) -> astropy.io.fits.BinTableHDU:
356 # We use a fixed-length string for the EXTNAME column; it might be
357 # better to use a variable-length array, but I have not been able to
358 # figure out how to get astropy to accept a string for the the
359 # character (TFORM='A') variant of that. And that's only better if the
360 # EXTNAMEs get super long, which is not likely (but maybe something to
361 # guard against).
362 max_name_size = max(len(hdu.header.get("EXTNAME", "")) for hdu in hdu_list)
363 index_hdu = astropy.io.fits.BinTableHDU.from_columns(
364 [
365 astropy.io.fits.Column("EXTNAME", f"A{max_name_size}"),
366 astropy.io.fits.Column("EXTVER", "J"),
367 astropy.io.fits.Column("XTENSION", "A8"),
368 astropy.io.fits.Column("ZIMAGE", "L"),
369 ]
370 + _HDUBytes.get_index_hdu_columns(),
371 nrows=len(hdu_list),
372 name="INDEX",
373 )
374 hdu: ExtensionHDU | astropy.io.fits.PrimaryHDU
375 for n, hdu in enumerate(hdu_list):
376 index_hdu.data[n]["EXTNAME"] = hdu.header.get("EXTNAME", "")
377 index_hdu.data[n]["EXTVER"] = hdu.header.get("EXTVER", 1)
378 index_hdu.data[n]["XTENSION"] = hdu.header.get("XTENSION", "IMAGE")
379 index_hdu.data[n]["ZIMAGE"] = hdu.header.get("ZIMAGE", False)
380 bytes = _HDUBytes.from_read_hdu(hdu)
381 bytes.update_index_row(index_hdu.data[n])
382 return index_hdu
385@dataclasses.dataclass
386class _HDUBytes:
387 """A struct that records the byte offsets into a FITS HDU."""
389 @classmethod
390 def from_write_hdu(cls, hdu: astropy.io.fits.PrimaryHDU | ExtensionHDU) -> Self:
391 """Construct from an Astropy HDU instance that has just been written.
393 Parameters
394 ----------
395 hdu
396 An Astropy HDU object.
398 Returns
399 -------
400 hdu_bytes
401 Struct with byte offsets.
403 Notes
404 -----
405 This method relies on internal Astropy attributes and does not work on
406 CompImageHDU objects.
407 """
408 # This is implemented by accessing private Astropy attributes because
409 # it turns out that's much more reliable than the public fileinfo()
410 # method, which seems to always return a dict with `None` entries or
411 # raise; it looks buggy, but docs are scarce enough that it's not clear
412 # what the right behavior is supposed to be.
413 if (header_address := getattr(hdu, "_header_offset", None)) is None:
414 raise RuntimeError("Failed to get Astropy's _header_offset.")
415 if (data_address := getattr(hdu, "_data_offset", None)) is None:
416 raise RuntimeError("Failed to get Astropy's _data_offset.")
417 if (data_size := getattr(hdu, "_data_size", None)) is None:
418 raise RuntimeError("Failed to get Astropy's _data_size.")
419 return cls(header_address, data_address, data_size)
421 @classmethod
422 def from_read_hdu(cls, hdu: astropy.io.fits.PrimaryHDU | ExtensionHDU) -> Self:
423 """Construct from an Astropy HDU instance that has just been read.
425 Parameters
426 ----------
427 hdu
428 An Astropy HDU object.
430 Returns
431 -------
432 hdu_bytes
433 Struct with byte offsets.
434 """
435 info = hdu.fileinfo()
436 header_address = info["hdrLoc"]
437 data_address = info["datLoc"]
438 data_size = info["datSpan"]
439 return cls(header_address, data_address, data_size)
441 @classmethod
442 def from_index_row(cls, row: np.void) -> Self:
443 """Construct from a row of the index HDU.
445 Parameters
446 ----------
447 row
448 A Numpy struct-like scalar.
450 Returns
451 -------
452 hdu_bytes
453 Struct with byte offsets.
454 """
455 return cls(
456 header_address=int(row["HDRADDR"]),
457 data_address=int(row["DATADDR"]),
458 data_size=int(row["DATSIZE"]),
459 )
461 @staticmethod
462 def get_index_hdu_columns() -> list[astropy.io.fits.Column]:
463 """Return the definitions of the columns this class gets and sets
464 from the index HDU.
466 Returns
467 -------
468 columns
469 A `list` of `astropy.io.fits.Column` objects that represent the
470 header address, data address, and data size.
471 """
472 return [
473 astropy.io.fits.Column("HDRADDR", "K"),
474 astropy.io.fits.Column("DATADDR", "K"),
475 astropy.io.fits.Column("DATSIZE", "K"),
476 ]
478 header_address: int
479 """Offset from the beginning of the start of the file to the header of this
480 HDU, in bytes.
481 """
483 data_address: int
484 """Offset from the beginning of the start of the file to the data section
485 of this HDU, in bytes.
486 """
488 data_size: int
489 """Size of the data section in bytes."""
491 @property
492 def header_size(self) -> int:
493 """Size of the header in bytes."""
494 return self.data_address - self.header_address
496 @property
497 def end_address(self) -> int:
498 """Offset in bytes from the start of the file to the end of the HDU."""
499 return self.data_address + self.data_size
501 @property
502 def size(self) -> int:
503 """Total size of this HDU in bytes."""
504 return self.data_size + self.data_address - self.header_address
506 def update_index_row(self, row: np.void) -> None:
507 """Set the values of a row of the index HDU from this strut.
509 Parameters
510 ----------
511 row
512 A Numpy struct-like scalar to modify in place.
513 """
514 row["HDRADDR"] = self.header_address
515 row["DATADDR"] = self.data_address
516 row["DATSIZE"] = self.data_size