Coverage for python/lsst/images/fits/_input_archive.py: 45%
246 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-19 02:55 -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 "DEFAULT_PAGE_SIZE",
16 "READ_CACHE_MAX_BYTES",
17 "FitsInputArchive",
18 "FitsOpaqueMetadata",
19)
21import io
22import os
23from collections.abc import Callable, Iterator
24from contextlib import contextmanager
25from functools import cached_property
26from types import EllipsisType
27from typing import IO, Any, Self
29import astropy.io.fits
30import astropy.table
31import fsspec
32import numpy as np
34from lsst.resources import ResourcePath, ResourcePathExpression
36from .._transforms import FrameSet
37from ..serialization import (
38 ArchiveInfo,
39 ArchiveReadError,
40 ArchiveTree,
41 ArrayReferenceModel,
42 InlineArrayModel,
43 InputArchive,
44 TableModel,
45 no_header_updates,
46 parameterize_tree,
47 tree_class_for_info,
48)
49from ..serialization._backends import _is_binary_stream
50from ..serialization._common import _ARCHIVE_READ_CONTEXT, _check_format_version
51from ._common import (
52 JSON_COLUMN,
53 JSON_EXTNAME,
54 ExtensionHDU,
55 ExtensionKey,
56 FitsOpaqueMetadata,
57 InvalidFitsArchiveError,
58 PointerModel,
59)
61_FITS_FORMAT_VERSION = 1
62"""Container layout version this release of `FitsInputArchive` understands."""
64DEFAULT_PAGE_SIZE = 2880 * 800
65"""Default fsspec read-block size for partial (remote) reads, in bytes.
67This is the single place to tune the block size for remote-store performance.
68On a buffered remote filesystem (e.g. GCS) each cache miss is one range
69request, so the block size trades round trips against over-fetch: a component
70read that touches scattered compressed tiles pulls one block per cluster of
71nearby tiles, rounded up to this size.
73The optimum depends on the access pattern. Larger blocks favor reads that
74touch most of the file (full planes, large cutouts); smaller blocks reduce
75wasted bytes for small scattered cutouts. ``2880 * 800`` (~2.3 MB, and a
76multiple of the 2880-byte FITS block) is a robust middle: across cutout sizes
77and full reads it stays within ~1.5x of the per-pattern optimum, whereas the
78historical 144 KB default was several times slower for all but the tiniest
79cutout. Raise it (e.g. ``2880 * 1600``) when whole-file or large-cutout reads
80dominate; lower it when many tiny cutouts across many files dominate.
82Local filesystems ignore this (their opener does no buffering), so it only
83affects remote stores.
84"""
86_READ_CACHE_TYPE = "blockcache"
87"""fsspec cache strategy for partial reads.
89``blockcache`` keeps a bounded set of fixed-size blocks (so memory stays
90capped) and reuses them across the multiple components of one file -- image,
91mask, variance and so on often share blocks -- unlike the default unbounded
92single-block ``readahead``.
93"""
95READ_CACHE_MAX_BYTES = 64 * 1024 * 1024
96"""Approximate memory budget for the partial-read block cache, per open file.
98The fsspec block cache evicts least-recently-used blocks once it holds more
99than ``maxblocks``; we derive ``maxblocks`` from this budget and the block
100size (`DEFAULT_PAGE_SIZE`) so the memory cap is expressed in bytes and stays
101fixed even when the block size is retuned. Measured benefit saturates at two
102cached blocks for the access patterns we care about, so this budget is purely
103headroom plus a guard against unbounded growth; it is far below fsspec's
104implicit default of ``32 * block_size``.
105"""
108class FitsInputArchive(InputArchive[PointerModel]):
109 """An implementation of the `.serialization.InputArchive` interface that
110 reads from FITS files.
112 Instances of this class should only be constructed via the `open`
113 context manager.
115 Parameters
116 ----------
117 stream
118 Open binary stream the archive reads from.
119 """
121 @classmethod
122 def get_basic_info(cls, path: ResourcePathExpression) -> ArchiveInfo:
123 """Read ``DATAMODL`` (schema URL) and ``FMTVER`` (container version)
124 from the primary header.
126 Every FITS file written by this package records the schema URL in
127 the ``DATAMODL`` card, so the schema can be identified without
128 reading the (potentially large) JSON tree HDU.
130 Parameters
131 ----------
132 path
133 Path to the archive to read.
134 """
135 with ResourcePath(path).open("rb") as stream:
136 primary = astropy.io.fits.PrimaryHDU.readfrom(stream)
137 header = primary.header
138 format_version = header.get("FMTVER")
139 schema_url = header.get("DATAMODL")
140 if not schema_url: 140 ↛ 141line 140 didn't jump to line 141 because the condition on line 140 was never true
141 raise ArchiveReadError(f"{path!r} is not an lsst.images FITS archive (no DATAMODL card).")
142 # DATAMODL was found, so this is one of our files and carries the
143 # layout stamp its writer emits first.
144 if format_version is None:
145 raise ArchiveReadError(f"{path!r} has a DATAMODL card but no FMTVER card.")
146 return ArchiveInfo.from_schema_url(schema_url, format_version=int(format_version))
148 @classmethod
149 @contextmanager
150 def open_tree(
151 cls,
152 path: ResourcePathExpression | IO[bytes],
153 *,
154 partial: bool = True,
155 **backend_kwargs: Any,
156 ) -> Iterator[tuple[Self, ArchiveTree, ArchiveInfo]]:
157 """Open the FITS file and yield ``(archive, tree, info)``.
159 Parameters
160 ----------
161 path
162 The file resource to open, or a seekable binary stream
163 containing the file's content.
164 partial
165 If `True` the file is opened without reading it all into memory.
166 **backend_kwargs
167 Optional parameters for this backend. Currently supports
168 ``page_size`` which can be used to override the default
169 page size (which can be overridden globally by modifying
170 `DEFAULT_PAGE_SIZE`).
171 """
172 page_size = backend_kwargs.pop("page_size", DEFAULT_PAGE_SIZE)
173 with cls.open(path, page_size=page_size, partial=partial) as archive:
174 info = archive.info
175 tree_cls = tree_class_for_info(info, path)
176 parameterized = parameterize_tree(tree_cls, PointerModel)
177 tree = archive.get_tree(parameterized)
178 yield archive, tree, info
180 def __init__(self, stream: IO[bytes]) -> None:
181 self._primary_hdu = astropy.io.fits.PrimaryHDU.readfrom(stream)
182 # Every file this class can open was written by FitsOutputArchive,
183 # which stamps FMTVER before writing anything else -- the container
184 # cards popped below have no defaults, so a file without them was
185 # never readable. A missing stamp is therefore a damaged file rather
186 # than an old one, and it is worth saying so before those pops turn it
187 # into a bare KeyError.
188 on_disk_fmtver: int | None = self._primary_hdu.header.pop("FMTVER", None)
189 if on_disk_fmtver is None:
190 raise ArchiveReadError("This is not an lsst.images FITS archive (no FMTVER card).")
191 # DATAMODL is informational only on read; the JSON tree's
192 # schema_version / min_read_version drive data-model checks. We
193 # capture it here as ArchiveInfo so callers (e.g. open_tree) can
194 # identify the schema from this open rather than reopening the file.
195 # A schema-less file can still be opened directly; only callers that
196 # need the schema (via the `info` property) require DATAMODL.
197 schema_url = self._primary_hdu.header.pop("DATAMODL", None)
198 self._info = (
199 ArchiveInfo.from_schema_url(schema_url, format_version=on_disk_fmtver) if schema_url else None
200 )
201 _check_format_version("fits", on_disk_fmtver, _FITS_FORMAT_VERSION)
202 # TODO: do some basic checks that the file format conforms to our
203 # expectations (e.g. primary HDU should have no data).
204 #
205 # Read and strip the addresses and sizes from the headers. We don't
206 # actually need the index address because we always want to read the
207 # JSON HDU, too, and the index HDU is always the next one (but this
208 # could change in the future).
209 json_address: int = self._primary_hdu.header.pop("JSONADDR")
210 json_size: int = self._primary_hdu.header.pop("JSONSIZE")
211 del self._primary_hdu.header["INDXADDR"]
212 index_size: int = self._primary_hdu.header.pop("INDXSIZE")
213 # Save the remaining primary header keys so we can propagate them on
214 # rewrite.
215 self._opaque_metadata = FitsOpaqueMetadata()
216 self._opaque_metadata.add_header(self._primary_hdu.header.copy(strip=True), name="", ver=1)
217 # Read the JSON and index HDUs from the end.
218 stream.seek(json_address)
219 tail_data = stream.read(json_size + index_size)
220 index_hdu = astropy.io.fits.BinTableHDU.fromstring(tail_data[json_size:])
221 # Initialize lazy readers for all of the regular HDUs and the JSON HDU.
222 self._readers = {
223 ExtensionKey.from_index_row(row): _ExtensionReader.from_index_row(row, stream)
224 for row in index_hdu.data
225 }
226 self._readers[ExtensionKey(JSON_COLUMN)] = _ExtensionReader.from_bytes(
227 astropy.io.fits.BinTableHDU, tail_data[:json_size]
228 )
229 # Make any empty dictionary to cache deserialized objects. Keys are
230 # the zero-indexed row in the JSON table.
231 self._deserialized_pointer_cache: dict[int, Any] = {}
233 @classmethod
234 @contextmanager
235 def open(
236 cls,
237 path: ResourcePathExpression | IO[bytes],
238 *,
239 page_size: int = DEFAULT_PAGE_SIZE,
240 partial: bool = False,
241 ) -> Iterator[Self]:
242 """Create an output archive that writes to the given file.
244 Parameters
245 ----------
246 path
247 File to read; convertible to `lsst.resources.ResourcePath`,
248 or a seekable binary stream containing the file's content.
249 For stream input ``page_size`` and ``partial`` are ignored:
250 the data is already in memory and needs no paging.
251 page_size
252 Size of the fsspec read block for partial (remote) reads, in
253 bytes; a multiple of the FITS block size (2880) is recommended.
254 Defaults to `DEFAULT_PAGE_SIZE`; see it for the tuning tradeoff.
255 partial
256 Whether we will be reading only some of the archive, or if memory
257 pressure forces us to read it only a little at a time. If `False`
258 (default), the entire raw file may be read into memory up front.
260 Returns
261 -------
262 `contextlib.AbstractContextManager` [`FitsInputArchive`]
263 A context manager that returns a `FitsInputArchive` when entered.
264 """
265 if _is_binary_stream(path):
266 yield cls(path)
267 return
268 path = ResourcePath(path)
269 stream: IO[bytes]
270 if not partial:
271 stream = io.BytesIO(path.read())
272 yield cls(stream)
273 else:
274 fs: fsspec.AbstractFileSystem
275 fs, fp = path.to_fsspec()
276 # Cap cached blocks from the byte budget so memory stays bounded as
277 # the block size is retuned; keep at least two so the shared
278 # header/index block survives between a file's components.
279 maxblocks = max(2, READ_CACHE_MAX_BYTES // page_size)
280 with fs.open(
281 fp,
282 block_size=page_size,
283 cache_type=_READ_CACHE_TYPE,
284 cache_options={"maxblocks": maxblocks},
285 ) as stream:
286 yield cls(stream)
288 @property
289 def info(self) -> ArchiveInfo:
290 """Schema/format info read from the primary header on open
291 (`.serialization.ArchiveInfo`).
292 """
293 if self._info is None:
294 raise ArchiveReadError("This is not an lsst.images FITS archive (no DATAMODL card).")
295 return self._info
297 def get_tree[T: ArchiveTree](self, model_type: type[T]) -> T:
298 """Read the JSON tree from the archive.
300 Parameters
301 ----------
302 model_type
303 A Pydantic model type to use to validate the JSON.
305 Returns
306 -------
307 T
308 The validated Pydantic model.
309 """
310 json_bytes = self._readers[ExtensionKey(JSON_EXTNAME)].data[0][JSON_COLUMN].tobytes()
311 return model_type.model_validate_json(json_bytes, context=_ARCHIVE_READ_CONTEXT)
313 def deserialize_pointer[U: ArchiveTree, V](
314 self,
315 pointer: PointerModel,
316 model_type: type[U],
317 deserializer: Callable[[U, InputArchive[PointerModel]], V],
318 ) -> V:
319 # Docstring inherited.
320 if (cached := self._deserialized_pointer_cache.get(pointer.row)) is not None:
321 return cached
322 if not isinstance(pointer.column.data, ArrayReferenceModel):
323 raise ArchiveReadError(f"Invalid pointer with inline array:\n{pointer.model_dump_json(indent=2)}")
324 _, reader = self._get_source_reader(pointer.column.data.source, is_table=True)
325 try:
326 json_bytes = reader.data[pointer.row][JSON_COLUMN].tobytes()
327 except Exception as err:
328 raise InvalidFitsArchiveError(
329 f"Failed to access the table cell referenced by {pointer.model_dump_json()}."
330 ) from err
331 result = deserializer(model_type.model_validate_json(json_bytes, context=_ARCHIVE_READ_CONTEXT), self)
332 self._deserialized_pointer_cache[pointer.row] = result
333 return result
335 def get_frame_set(self, ref: PointerModel) -> FrameSet:
336 try:
337 result = self._deserialized_pointer_cache[ref.row]
338 except KeyError:
339 raise AssertionError(
340 f"Frame set at {ref.model_dump_json(indent=2)} must be deserialized "
341 "before any dependent transform can be."
342 ) from None
343 if not isinstance(result, FrameSet):
344 raise InvalidFitsArchiveError(f"Expected a FrameSet instance at {ref.model_dump_json(indent=2)}.")
345 return result
347 def get_array(
348 self,
349 model: ArrayReferenceModel | InlineArrayModel,
350 *,
351 slices: tuple[slice, ...] | EllipsisType = ...,
352 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
353 ) -> np.ndarray:
354 if not isinstance(model, ArrayReferenceModel): 354 ↛ 355line 354 didn't jump to line 355 because the condition on line 354 was never true
355 raise ArchiveReadError("Inline array found where a reference array was expected.")
356 key, reader = self._get_source_reader(model.source, is_table=False)
357 if slices is not ...:
358 array = reader.section[slices]
359 else:
360 array = reader.data
361 if key not in self._opaque_metadata.headers: 361 ↛ 365line 361 didn't jump to line 365 because the condition on line 361 was always true
362 opaque_header = reader.header.copy(strip=True)
363 strip_header(opaque_header)
364 self._opaque_metadata.add_header(opaque_header, key=key)
365 return array
367 def get_table(
368 self,
369 model: TableModel,
370 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
371 ) -> astropy.table.Table:
372 # Docstring inherited.
373 array = self.get_structured_array(model, strip_header)
374 table = astropy.table.Table(array)
375 for c in model.columns:
376 c.update_table(table)
377 return table
379 def get_structured_array(
380 self,
381 model: TableModel,
382 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
383 ) -> np.ndarray:
384 # Docstring inherited.
385 if not isinstance(model.columns[0].data, ArrayReferenceModel):
386 raise ArchiveReadError("Inline array found where a reference array was expected.")
387 # All columns should have the same data.source; just use the first.
388 key, reader = self._get_source_reader(model.columns[0].data.source, is_table=True)
389 if key not in self._opaque_metadata.headers:
390 opaque_header = reader.header.copy(strip=True)
391 strip_header(opaque_header)
392 self._opaque_metadata.add_header(opaque_header, key=key)
393 return reader.hdu.data
395 def get_opaque_metadata(self) -> FitsOpaqueMetadata:
396 # Docstring inherited.
397 return self._opaque_metadata
399 def _get_source_reader(self, source: str | int, is_table: bool) -> tuple[ExtensionKey, _ExtensionReader]:
400 """Get a reader for the extension referenced by a serialiation model's
401 ``source`` field.
403 Parameters
404 ----------
405 source
406 A ``source`` field of the form ``fits:${hdu}`` or
407 ``fits:${hdu}[${col}]``.
408 is_table
409 Whether the source should be for a table HDU.
411 Returns
412 -------
413 key
414 Identifier pair for the HDU (EXTNAME, EXTVER).
415 reader
416 A reader object for the extension.
417 """
418 if not isinstance(source, str): 418 ↛ 419line 418 didn't jump to line 419 because the condition on line 418 was never true
419 raise InvalidFitsArchiveError(f"Reference with source={source!r} is not a string.")
420 if not source.startswith("fits:"): 420 ↛ 421line 420 didn't jump to line 421 because the condition on line 420 was never true
421 raise InvalidFitsArchiveError(f"Reference with source={source!r} does not start with 'fits:'.")
422 key = ExtensionKey.from_str(source)
423 try:
424 reader = self._readers[key]
425 except KeyError:
426 raise InvalidFitsArchiveError(f"Unrecognized source value {key}.") from None
427 if is_table and not reader.is_table: 427 ↛ 428line 427 didn't jump to line 428 because the condition on line 427 was never true
428 raise InvalidFitsArchiveError(
429 f"Extension with source={key} was expected to be be a binary table, not an image."
430 )
431 elif not is_table and reader.is_table: 431 ↛ 432line 431 didn't jump to line 432 because the condition on line 431 was never true
432 raise InvalidFitsArchiveError(
433 f"Extension with source={key} was expected to be be an image, not a binary table."
434 )
435 return key, reader
438class _ExtensionReader:
439 """A lazy-load reader for a single extension HDU.
441 Parameters
442 ----------
443 hdu_cls
444 The type of the astropy HDU instance to construct.
445 stream
446 The file-like object to read from.
447 """
449 def __init__(self, hdu_cls: type[ExtensionHDU], stream: IO[bytes]) -> None:
450 self._hdu_cls = hdu_cls
451 self._stream = stream
453 @classmethod
454 def from_index_row(cls, index_row: np.void, stream: IO[bytes]) -> _ExtensionReader:
455 """Construct from a row of the binary table index HDU.
457 Parameters
458 ----------
459 index_row
460 A record array row from the index HDU.
461 stream
462 The file-like object being used to read the full FITS file.
464 Returns
465 -------
466 reader
467 A reader object for the extension.
468 """
469 match index_row["XTENSION"].strip():
470 case "IMAGE":
471 hdu_cls = astropy.io.fits.ImageHDU
472 case "BINTABLE": 472 ↛ 477line 472 didn't jump to line 477 because the pattern on line 472 always matched
473 if index_row["ZIMAGE"]:
474 hdu_cls = astropy.io.fits.CompImageHDU
475 else:
476 hdu_cls = astropy.io.fits.BinTableHDU
477 case other:
478 raise AssertionError(f"Unsupported HDU type {other!r}.")
479 return _ExtensionReader(
480 hdu_cls,
481 _RangeStreamProxy(
482 stream,
483 start=int(index_row["HDRADDR"]),
484 ),
485 )
487 @classmethod
488 def from_bytes(cls, hdu_cls: type[ExtensionHDU], data: bytes) -> _ExtensionReader:
489 """Construct from already-read `bytes`.
491 Parameters
492 ----------
493 hdu_cls
494 The HDU type to instantiate.
495 data
496 Raw data for the HDU.
498 Returns
499 -------
500 reader
501 A reader object for extension.
502 """
503 return _ExtensionReader(hdu_cls, io.BytesIO(data))
505 @property
506 def is_table(self) -> bool:
507 """Whether this is logically a table HDU.
509 This is `False` for compressed image HDUs, even though they are
510 represented in FITS as a binary table.
511 """
512 return issubclass(self._hdu_cls, astropy.io.fits.BinTableHDU) and not issubclass(
513 self._hdu_cls, astropy.io.fits.CompImageHDU
514 )
516 @cached_property
517 def hdu(self) -> ExtensionHDU:
518 """The Astropy HDU object."""
519 self._stream.seek(0)
520 if self._hdu_cls is astropy.io.fits.CompImageHDU:
521 # CompImageHDU.readfrom doesn't work; we need to make a minimal
522 # example and report it upstream. Happily this workaround does
523 # work.
524 bintable_hdu = astropy.io.fits.BinTableHDU.readfrom(self._stream, memmap=False, cache=False)
525 return self._hdu_cls(bintable=bintable_hdu)
526 else:
527 return self._hdu_cls.readfrom(self._stream, memmap=False, cache=False, uint=True)
529 @property
530 def header(self) -> astropy.io.fits.Header:
531 """The header of the HDU."""
532 return self.hdu.header
534 @property
535 def data(self) -> np.ndarray:
536 """The data for the HDU."""
537 return self.hdu.data
539 @property
540 def section(self) -> astropy.io.fits.Section | astropy.io.fits.CompImageSection:
541 """An Astropy expression object that reads a subset of the data when
542 sliced.
543 """
544 return self.hdu.section
547class _RangeStreamProxy(IO[bytes]):
548 """A readable IO proxy object that makes the beginning of the file appear
549 at a custom position.
551 Parameters
552 ----------
553 base
554 Underlying readable, seekable buffer to proxy.
555 start
556 Offset into the base stream that will be considered the start of the
557 proxy stream.
559 Notes
560 -----
561 This class exists because Astropy doesn't seem to provide a way to read a
562 single HDU that starts at the current seek position of a file-like object.
563 It does provide ``readfrom`` methods on its HDU objects that take a
564 file-like object, but these assume (possibly unintentionally; it only
565 happens when Astropy is trying to see whether the file was opened for
566 appending) that ``seek(0)`` will set the file-like object to the start of
567 the HDU.
568 """
570 def __init__(self, base: IO[bytes], start: int) -> None:
571 self._base = base
572 self._start = start
574 @property
575 def mode(self) -> str:
576 return "rb"
578 def __enter__(self) -> Self:
579 raise AssertionError("This proxy should not be used as a context manager.")
581 def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
582 raise AssertionError("This proxy should not be used as a context manager.")
584 def __iter__(self) -> Iterator[bytes]:
585 return self._base.__iter__()
587 def __next__(self) -> bytes:
588 return self._base.__next__()
590 def close(self) -> None:
591 raise AssertionError("This proxy should not ever be closed.")
593 @property
594 def closed(self) -> bool:
595 return False
597 def fileno(self) -> int:
598 raise OSError()
600 def flush(self) -> None:
601 pass
603 def isatty(self) -> bool:
604 return False
606 def read(self, n: int = -1, /) -> bytes:
607 result = self._base.read(n)
608 return result
610 def readable(self) -> bool:
611 return True
613 def readline(self, limit: int = -1, /) -> bytes:
614 return self._base.readline(limit)
616 def readlines(self, hint: int = -1, /) -> list[bytes]:
617 return self._base.readlines(hint)
619 def seek(self, offset: int, whence: int = 0) -> int:
620 match whence:
621 case os.SEEK_SET:
622 return self._base.seek(offset + self._start, os.SEEK_SET) - self._start
623 case os.SEEK_CUR: 623 ↛ 624line 623 didn't jump to line 624 because the pattern on line 623 never matched
624 return self._base.seek(offset, os.SEEK_CUR) - self._start
625 case os.SEEK_END: 625 ↛ 627line 625 didn't jump to line 627 because the pattern on line 625 always matched
626 return self._base.seek(offset, os.SEEK_END) - self._start
627 raise TypeError(f"Invalid value for 'whence': {whence}.")
629 def seekable(self) -> bool:
630 return True
632 def tell(self) -> int:
633 return self._base.tell() - self._start
635 def truncate(self, size: int | None = None, /) -> int:
636 raise OSError()
638 def writable(self) -> bool:
639 return False
641 def write(self, arg: Any, /) -> int:
642 raise OSError()
644 def writelines(self, arg: Any, /) -> None:
645 raise OSError()