Coverage for python/lsst/images/fits/_output_archive.py: 20%
224 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-12 00:48 -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
17import warnings
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 MetadataValue,
33 NestedOutputArchive,
34 NumberType,
35 OutputArchive,
36 TableColumnModel,
37 TableModel,
38 no_header_updates,
39)
40from ._common import (
41 JSON_COLUMN,
42 JSON_EXTNAME,
43 ExtensionHDU,
44 ExtensionKey,
45 FitsCompressionOptions,
46 FitsOpaqueMetadata,
47 PointerModel,
48)
50_FITS_FORMAT_VERSION = 1
51"""Container layout version for files written by `FitsOutputArchive`.
53Bumps when the on-disk FITS layout (HDU placement, INDX/JSON keyword schema)
54changes. Independent of any data-model ``SCHEMA_VERSION``.
55"""
58def write(
59 obj: Any,
60 path: str,
61 compression_options: Mapping[str, FitsCompressionOptions | None] | None = None,
62 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
63 compression_seed: int | None = None,
64 metadata: dict[str, MetadataValue] | None = None,
65 butler_info: ButlerInfo | None = None,
66) -> Any:
67 """Write an object with a ``serialize`` method to a FITS file.
69 Parameters
70 ----------
71 path
72 Name of the file to write to. Must not already exist.
73 compression_options
74 Options for how to compress the FITS file, keyed by the name of
75 the attribute (with JSON pointer ``/`` separators for nested
76 attributes).
77 update_header
78 A callback that will be given the primary HDU FITS header and an
79 opportunity to modify it.
80 compression_seed
81 A FITS tile compression seed to use whenever the configured
82 compression seed is `None` or (for backwards compatibility) ``0``.
83 This value is then incremented every time it is used.
84 metadata
85 Additional metadata to save with the object. This will override any
86 flexible metadata carried by the object itself with the same keys.
87 butler_info
88 Butler information to store in the file.
90 Returns
91 -------
92 `.serialization.ArchiveTree`
93 The serialized representation of the object.
94 """
95 opaque_metadata = getattr(obj, "_opaque_metadata", None)
96 name = getattr(obj, "_archive_default_name", None)
97 with FitsOutputArchive.open(
98 path,
99 compression_options=compression_options,
100 opaque_metadata=opaque_metadata,
101 update_header=update_header,
102 compression_seed=compression_seed,
103 ) as archive:
104 tree = archive.serialize_direct(name, obj.serialize) if name is not None else obj.serialize(archive)
105 if metadata is not None:
106 tree.metadata.update(metadata)
107 if butler_info is not None:
108 tree.butler_info = butler_info
109 archive.add_tree(tree)
110 return tree
113class FitsOutputArchive(OutputArchive[PointerModel]):
114 """An implementation of the `.serialization.OutputArchive` interface that
115 writes to FITS files.
117 Instances of this class should only be constructed via the `open`
118 context manager.
119 """
121 def __init__(
122 self,
123 hdu_list: astropy.io.fits.HDUList,
124 compression_options: Mapping[str, FitsCompressionOptions | None] | None = None,
125 opaque_metadata: Any = None,
126 compression_seed: int | None = None,
127 ):
128 super().__init__()
129 # JSON blobs for objects we've saved as pointers:
130 self._pointer_targets: list[bytes] = []
131 # Mapping from user provided key (e.g. id(some object)) to a table
132 # pointer to where we actually saved it:
133 self._pointers_by_key: dict[Hashable, PointerModel] = {}
134 self._hdu_list = hdu_list
135 self._primary_hdu = astropy.io.fits.PrimaryHDU()
136 self._primary_hdu.header.set("FMTVER", _FITS_FORMAT_VERSION, "FITS container layout version.")
137 self._primary_hdu.header.set("INDXADDR", 0, "Offset in bytes to the HDU index.")
138 self._primary_hdu.header.set("INDXSIZE", 0, "Size of the HDU index.")
139 self._primary_hdu.header.set("JSONADDR", 0, "Offset in bytes to the JSON tree HDU.")
140 self._primary_hdu.header.set("JSONSIZE", 0, "Size of the JSON tree HDU.")
141 self._compression_options = dict(compression_options) if compression_options is not None else {}
142 self._compression_seed = compression_seed
143 self._opaque_metadata = (
144 opaque_metadata if isinstance(opaque_metadata, FitsOpaqueMetadata) else FitsOpaqueMetadata()
145 )
146 if (opaque_primary_header := self._opaque_metadata.headers.get(ExtensionKey())) is not None:
147 self._primary_hdu.header.extend(opaque_primary_header)
148 self._hdu_list.append(self._primary_hdu)
149 self._json_hdu_added: bool = False
150 self._frame_sets: list[tuple[FrameSet, PointerModel]] = []
152 @classmethod
153 @contextmanager
154 def open(
155 cls,
156 filename: str,
157 compression_options: Mapping[str, FitsCompressionOptions | None] | None = None,
158 opaque_metadata: Any = None,
159 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
160 compression_seed: int | None = None,
161 ) -> Iterator[Self]:
162 """Create an output archive that writes to the given file.
164 Parameters
165 ----------
166 filename
167 Name of the file to write to. Must not already exist.
168 compression_options
169 Options for how to compress the FITS file, keyed by the name of
170 the attribute (with JSON pointer ``/`` separators for nested
171 attributes).
172 opaque_metadata
173 Metadata read from an input archive along with the object being
174 written now. Ignored if the metadata is not from a FITS archive.
175 update_header
176 A callback that will be given the primary HDU FITS header and an
177 opportunity to modify it.
178 compression_seed
179 A FITS tile compression seed to use whenever the configured
180 compression seed is `None` or (for backwards compatibility) ``0``.
181 This value is then incremented every time it is used.
183 Returns
184 -------
185 `contextlib.AbstractContextManager` [`FitsOutputArchive`]
186 A context manager that returns a `FitsOutputArchive` when entered.
187 """
188 with astropy.io.fits.open(filename, mode="append") as hdu_list:
189 if hdu_list:
190 raise OSError(f"File {filename!r} already exists.")
191 archive = cls(hdu_list, compression_options, opaque_metadata, compression_seed=compression_seed)
192 update_header(hdu_list[0].header)
193 yield archive
194 if not archive._json_hdu_added:
195 raise RuntimeError("Write context exited without 'add_tree' being called.")
196 # If a header card has a long string that required CONTINUE,
197 # Astropy will truncate the comment and warn without reporting
198 # what the offending card is. But it doesn't look at its own
199 # output_verify kwarg when doing that, and it doesn't actually
200 # trigger if you try to format the header cards one at a time!
201 # So we have no choice but to silence the warnings manually, until
202 # we can get Astropy fixed.
203 with warnings.catch_warnings():
204 warnings.filterwarnings(
205 "ignore",
206 message="comment will be truncated",
207 category=astropy.io.fits.verify.VerifyWarning,
208 )
209 hdu_list.flush()
210 # This multi-open dance is necessary to get Astropy to tell us the
211 # byte addresses of the HDUs. Hopefully we can get an upstream change
212 # make this unnecessary at some point.
213 with astropy.io.fits.open(filename, mode="readonly", disable_image_compression=True) as hdu_list:
214 index_hdu = cls._make_index_table(hdu_list)
215 with astropy.io.fits.open(filename, mode="append") as hdu_list:
216 hdu_list.append(index_hdu)
217 json_bytes = _HDUBytes.from_index_row(index_hdu.data[-1])
218 index_bytes = _HDUBytes.from_write_hdu(index_hdu)
219 # Update the primary HDU with the address and size of the index and
220 # JSON HDUs, and rewrite just that. We do this write manually, since
221 # astropy's docs on its 'update' mode are scarce and it's not obvious
222 # whether we can guarantee it won't rewrite the whole file if we edit
223 # the primary header.
224 archive._primary_hdu.header["INDXADDR"] = index_bytes.header_address
225 archive._primary_hdu.header["INDXSIZE"] = index_bytes.size
226 archive._primary_hdu.header["JSONADDR"] = json_bytes.header_address
227 archive._primary_hdu.header["JSONSIZE"] = json_bytes.size
228 with open(filename, "r+b") as stream:
229 stream.write(archive._primary_hdu.header.tostring().encode())
231 def serialize_direct[T: pydantic.BaseModel | None](
232 self, name: str, serializer: Callable[[OutputArchive[PointerModel]], T]
233 ) -> T:
234 nested = NestedOutputArchive[PointerModel](name, self)
235 return serializer(nested)
237 def serialize_pointer[T: ArchiveTree](
238 self, name: str, serializer: Callable[[OutputArchive[PointerModel]], T], key: Hashable
239 ) -> PointerModel:
240 if (pointer := self._pointers_by_key.get(key)) is not None:
241 return pointer
242 # Rooting the nested archive at "" follows the NestedOutputArchive
243 # convention that root paths are absolute, but serialize_direct roots
244 # at the bare name, so an array written by a pointer target would get
245 # a "/"-prefixed EXTNAME (e.g. "/DATA"). No serializer currently
246 # writes arrays inside pointer targets; if one ever does, consider
247 # rooting at ``name`` instead, as the NDF backend does.
248 model = self.serialize_direct("", serializer)
249 json_bytes = model.model_dump_json().encode()
250 self._pointer_targets.append(json_bytes)
251 pointer = PointerModel(
252 column=TableColumnModel(
253 name=JSON_COLUMN,
254 data=ArrayReferenceModel(
255 source=f"fits:{JSON_EXTNAME}[1]",
256 shape=[len(json_bytes)],
257 datatype=NumberType.uint8,
258 ),
259 ),
260 row=len(self._pointer_targets),
261 )
262 self._pointers_by_key[key] = pointer
263 return pointer
265 def serialize_frame_set[T: ArchiveTree](
266 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
267 ) -> PointerModel:
268 # Docstring inherited.
269 pointer = self.serialize_pointer(name, serializer, key)
270 self._frame_sets.append((frame_set, pointer))
271 return pointer
273 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, PointerModel]]:
274 return iter(self._frame_sets)
276 def add_array(
277 self,
278 array: np.ndarray,
279 *,
280 name: str | None = None,
281 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
282 tile_shape: tuple[int, ...] | None = None,
283 options_name: str | None = None,
284 ) -> ArrayReferenceModel:
285 if name is None:
286 raise RuntimeError("Cannot save array with name=None unless it is nested.")
287 name, version = self._register_name(name)
288 extname = name.upper()
289 hdu = self._opaque_metadata.maybe_use_precompressed(extname)
290 if hdu is None:
291 if options_name is None:
292 options_name = name
293 if (compression_options := self._get_compression_options(options_name, tile_shape)) is not None:
294 hdu = compression_options.make_hdu(array, name=extname)
295 else:
296 hdu = astropy.io.fits.ImageHDU(array, name=extname)
297 key = self._add_hdu(hdu, version, update_header)
298 return ArrayReferenceModel(
299 source=str(key),
300 shape=list(array.shape),
301 datatype=NumberType.from_numpy(array.dtype),
302 )
304 def add_table(
305 self,
306 table: astropy.table.Table,
307 *,
308 name: str | None = None,
309 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
310 ) -> TableModel:
311 if name is None:
312 raise RuntimeError("Cannot save table with name=None unless it is nested.")
313 name, version = self._register_name(name)
314 extname = name.upper()
315 hdu: astropy.io.fits.BinTableHDU = astropy.io.fits.table_to_hdu(table, name=extname)
316 # Extract column information directly from the input array, not the
317 # data in the binary table HDU, because we want to assume as little as
318 # possible about where Astropy does uint -> TZERO stuff.
319 columns = TableColumnModel.from_table(table)
320 key = self._add_hdu(hdu, version, update_header)
321 for n, c in enumerate(columns, start=1):
322 assert isinstance(c.data, ArrayReferenceModel)
323 c.data.source = f"{key}[{n}]"
324 return TableModel(columns=columns, meta=table.meta)
326 def add_structured_array(
327 self,
328 array: np.ndarray,
329 *,
330 name: str | None = None,
331 units: Mapping[str, astropy.units.Unit] | None = None,
332 descriptions: Mapping[str, str] | None = None,
333 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
334 ) -> TableModel:
335 if name is None:
336 raise RuntimeError("Cannot save structured array with name=None unless it is nested.")
337 name, version = self._register_name(name)
338 extname = name.upper()
339 # Extract column information directly from the input array, not the
340 # data in the binary table HDU, because we want to assume as little as
341 # possible about where Astropy does uint -> TZERO stuff.
342 columns = TableColumnModel.from_record_dtype(array.dtype)
343 hdu = astropy.io.fits.BinTableHDU(array, name=extname)
344 if units is not None:
345 for c in columns:
346 c.unit = units.get(c.name)
347 if descriptions is not None:
348 for c in columns:
349 c.description = descriptions.get(c.name, "")
350 key = self._add_hdu(hdu, version, update_header)
351 for n, c in enumerate(columns, start=1):
352 assert isinstance(c.data, ArrayReferenceModel)
353 c.data.source = f"{key}[{n}]"
354 return TableModel(columns=columns)
356 def _add_hdu(
357 self,
358 hdu: astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU | astropy.io.fits.BinTableHDU,
359 version: int,
360 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
361 ) -> ExtensionKey:
362 key = ExtensionKey(hdu.name, version)
363 key.check()
364 if version > 1:
365 hdu.header["EXTVER"] = version
366 update_header(hdu.header)
367 if (opaque_headers := self._opaque_metadata.headers.get(key)) is not None:
368 hdu.header.extend(opaque_headers)
369 self._hdu_list.append(hdu)
370 return key
372 def add_tree(self, tree: ArchiveTree) -> None:
373 """Write the JSON tree to the archive.
375 This method must be called exactly once, just before the `open` context
376 is exited.
378 Parameters
379 ----------
380 tree
381 Pydantic model that represents the tree.
382 """
383 self._primary_hdu.header.set("DATAMODL", tree.schema_url, "Schema URL.")
384 json_hdu = astropy.io.fits.BinTableHDU.from_columns(
385 [astropy.io.fits.Column(JSON_COLUMN, "PB")],
386 nrows=len(self._pointer_targets) + 1,
387 name=JSON_EXTNAME,
388 )
389 json_hdu.data[0][JSON_COLUMN] = np.frombuffer(tree.model_dump_json().encode(), dtype=np.byte)
390 for n, json_target_data in enumerate(self._pointer_targets):
391 json_hdu.data[n + 1][JSON_COLUMN] = np.frombuffer(json_target_data, dtype=np.byte)
392 self._hdu_list.append(json_hdu)
393 self._json_hdu_added = True
395 def _get_compression_options(
396 self, name: str, tile_shape: tuple[int, ...] | None
397 ) -> FitsCompressionOptions | None:
398 result = self._compression_options.get(name, FitsCompressionOptions.DEFAULT)
399 if result is None:
400 return result
401 if tile_shape is not None and result.tile_shape is None:
402 result = result.model_copy(update={"tile_shape": tile_shape})
403 if result.quantization is None:
404 return result
405 if self._compression_seed is not None and not result.quantization.seed:
406 result = result.model_copy(
407 update={
408 "quantization": result.quantization.model_copy(update={"seed": self._compression_seed})
409 }
410 )
411 self._compression_seed += 1
412 if self._compression_seed > 10000:
413 self._compression_seed = 1
414 # MyPy can tell that result.quantization is not None in the 'if', but
415 # forgets that by this 'else':
416 elif result.quantization.seed is None: # type: ignore[union-attr]
417 raise RuntimeError("No quantization seed provided.")
418 return result
420 @staticmethod
421 def _make_index_table(hdu_list: astropy.io.fits.HDUList) -> astropy.io.fits.BinTableHDU:
422 # We use a fixed-length string for the EXTNAME column; it might be
423 # better to use a variable-length array, but I have not been able to
424 # figure out how to get astropy to accept a string for the the
425 # character (TFORM='A') variant of that. And that's only better if the
426 # EXTNAMEs get super long, which is not likely (but maybe something to
427 # guard against).
428 max_name_size = max(len(hdu.header.get("EXTNAME", "")) for hdu in hdu_list)
429 index_hdu = astropy.io.fits.BinTableHDU.from_columns(
430 [
431 astropy.io.fits.Column("EXTNAME", f"A{max_name_size}"),
432 astropy.io.fits.Column("EXTVER", "J"),
433 astropy.io.fits.Column("XTENSION", "A8"),
434 astropy.io.fits.Column("ZIMAGE", "L"),
435 ]
436 + _HDUBytes.get_index_hdu_columns(),
437 nrows=len(hdu_list),
438 name="INDEX",
439 )
440 hdu: ExtensionHDU | astropy.io.fits.PrimaryHDU
441 for n, hdu in enumerate(hdu_list):
442 index_hdu.data[n]["EXTNAME"] = hdu.header.get("EXTNAME", "")
443 index_hdu.data[n]["EXTVER"] = hdu.header.get("EXTVER", 1)
444 index_hdu.data[n]["XTENSION"] = hdu.header.get("XTENSION", "IMAGE")
445 index_hdu.data[n]["ZIMAGE"] = hdu.header.get("ZIMAGE", False)
446 bytes = _HDUBytes.from_read_hdu(hdu)
447 bytes.update_index_row(index_hdu.data[n])
448 return index_hdu
451@dataclasses.dataclass
452class _HDUBytes:
453 """A struct that records the byte offsets into a FITS HDU."""
455 @classmethod
456 def from_write_hdu(cls, hdu: astropy.io.fits.PrimaryHDU | ExtensionHDU) -> Self:
457 """Construct from an Astropy HDU instance that has just been written.
459 Parameters
460 ----------
461 hdu
462 An Astropy HDU object.
464 Returns
465 -------
466 hdu_bytes
467 Struct with byte offsets.
469 Notes
470 -----
471 This method relies on internal Astropy attributes and does not work on
472 CompImageHDU objects.
473 """
474 # This is implemented by accessing private Astropy attributes because
475 # it turns out that's much more reliable than the public fileinfo()
476 # method, which seems to always return a dict with `None` entries or
477 # raise; it looks buggy, but docs are scarce enough that it's not clear
478 # what the right behavior is supposed to be.
479 if (header_address := getattr(hdu, "_header_offset", None)) is None:
480 raise RuntimeError("Failed to get Astropy's _header_offset.")
481 if (data_address := getattr(hdu, "_data_offset", None)) is None:
482 raise RuntimeError("Failed to get Astropy's _data_offset.")
483 if (data_size := getattr(hdu, "_data_size", None)) is None:
484 raise RuntimeError("Failed to get Astropy's _data_size.")
485 return cls(header_address, data_address, data_size)
487 @classmethod
488 def from_read_hdu(cls, hdu: astropy.io.fits.PrimaryHDU | ExtensionHDU) -> Self:
489 """Construct from an Astropy HDU instance that has just been read.
491 Parameters
492 ----------
493 hdu
494 An Astropy HDU object.
496 Returns
497 -------
498 hdu_bytes
499 Struct with byte offsets.
500 """
501 info = hdu.fileinfo()
502 header_address = info["hdrLoc"]
503 data_address = info["datLoc"]
504 data_size = info["datSpan"]
505 return cls(header_address, data_address, data_size)
507 @classmethod
508 def from_index_row(cls, row: np.void) -> Self:
509 """Construct from a row of the index HDU.
511 Parameters
512 ----------
513 row
514 A Numpy struct-like scalar.
516 Returns
517 -------
518 hdu_bytes
519 Struct with byte offsets.
520 """
521 return cls(
522 header_address=int(row["HDRADDR"]),
523 data_address=int(row["DATADDR"]),
524 data_size=int(row["DATSIZE"]),
525 )
527 @staticmethod
528 def get_index_hdu_columns() -> list[astropy.io.fits.Column]:
529 """Return the definitions of the columns this class gets and sets
530 from the index HDU.
532 Returns
533 -------
534 columns
535 A `list` of `astropy.io.fits.Column` objects that represent the
536 header address, data address, and data size.
537 """
538 return [
539 astropy.io.fits.Column("HDRADDR", "K"),
540 astropy.io.fits.Column("DATADDR", "K"),
541 astropy.io.fits.Column("DATSIZE", "K"),
542 ]
544 header_address: int
545 """Offset from the beginning of the start of the file to the header of this
546 HDU, in bytes.
547 """
549 data_address: int
550 """Offset from the beginning of the start of the file to the data section
551 of this HDU, in bytes.
552 """
554 data_size: int
555 """Size of the data section in bytes."""
557 @property
558 def header_size(self) -> int:
559 """Size of the header in bytes."""
560 return self.data_address - self.header_address
562 @property
563 def end_address(self) -> int:
564 """Offset in bytes from the start of the file to the end of the HDU."""
565 return self.data_address + self.data_size
567 @property
568 def size(self) -> int:
569 """Total size of this HDU in bytes."""
570 return self.data_size + self.data_address - self.header_address
572 def update_index_row(self, row: np.void) -> None:
573 """Set the values of a row of the index HDU from this strut.
575 Parameters
576 ----------
577 row
578 A Numpy struct-like scalar to modify in place.
579 """
580 row["HDRADDR"] = self.header_address
581 row["DATADDR"] = self.data_address
582 row["DATSIZE"] = self.data_size