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