Coverage for python/lsst/images/fits/_input_archive.py: 73%
242 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 02:27 -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 _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 = int(header.get("FMTVER", 1))
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 return ArchiveInfo.from_schema_url(schema_url, format_version=format_version)
144 @classmethod
145 @contextmanager
146 def open_tree(
147 cls,
148 path: ResourcePathExpression | IO[bytes],
149 *,
150 partial: bool = True,
151 **backend_kwargs: Any,
152 ) -> Iterator[tuple[Self, ArchiveTree, ArchiveInfo]]:
153 """Open the FITS file and yield ``(archive, tree, info)``.
155 Parameters
156 ----------
157 path
158 The file resource to open, or a seekable binary stream
159 containing the file's content.
160 partial
161 If `True` the file is opened without reading it all into memory.
162 **backend_kwargs
163 Optional parameters for this backend. Currently supports
164 ``page_size`` which can be used to override the default
165 page size (which can be overridden globally by modifying
166 `DEFAULT_PAGE_SIZE`).
167 """
168 page_size = backend_kwargs.pop("page_size", DEFAULT_PAGE_SIZE)
169 with cls.open(path, page_size=page_size, partial=partial) as archive:
170 info = archive.info
171 tree_cls = tree_class_for_info(info, path)
172 parameterized = parameterize_tree(tree_cls, PointerModel)
173 tree = archive.get_tree(parameterized)
174 yield archive, tree, info
176 def __init__(self, stream: IO[bytes]) -> None:
177 self._primary_hdu = astropy.io.fits.PrimaryHDU.readfrom(stream)
178 on_disk_fmtver: int = self._primary_hdu.header.pop("FMTVER", 1)
179 # DATAMODL is informational only on read; the JSON tree's
180 # schema_version / min_read_version drive data-model checks. We
181 # capture it here as ArchiveInfo so callers (e.g. open_tree) can
182 # identify the schema from this open rather than reopening the file.
183 # A schema-less file can still be opened directly; only callers that
184 # need the schema (via the `info` property) require DATAMODL.
185 schema_url = self._primary_hdu.header.pop("DATAMODL", None)
186 self._info = (
187 ArchiveInfo.from_schema_url(schema_url, format_version=on_disk_fmtver) if schema_url else None
188 )
189 _check_format_version("fits", on_disk_fmtver, _FITS_FORMAT_VERSION)
190 # TODO: do some basic checks that the file format conforms to our
191 # expectations (e.g. primary HDU should have no data).
192 #
193 # Read and strip the addresses and sizes from the headers. We don't
194 # actually need the index address because we always want to read the
195 # JSON HDU, too, and the index HDU is always the next one (but this
196 # could change in the future).
197 json_address: int = self._primary_hdu.header.pop("JSONADDR")
198 json_size: int = self._primary_hdu.header.pop("JSONSIZE")
199 del self._primary_hdu.header["INDXADDR"]
200 index_size: int = self._primary_hdu.header.pop("INDXSIZE")
201 # Save the remaining primary header keys so we can propagate them on
202 # rewrite.
203 self._opaque_metadata = FitsOpaqueMetadata()
204 self._opaque_metadata.add_header(self._primary_hdu.header.copy(strip=True), name="", ver=1)
205 # Read the JSON and index HDUs from the end.
206 stream.seek(json_address)
207 tail_data = stream.read(json_size + index_size)
208 index_hdu = astropy.io.fits.BinTableHDU.fromstring(tail_data[json_size:])
209 # Initialize lazy readers for all of the regular HDUs and the JSON HDU.
210 self._readers = {
211 ExtensionKey.from_index_row(row): _ExtensionReader.from_index_row(row, stream)
212 for row in index_hdu.data
213 }
214 self._readers[ExtensionKey(JSON_COLUMN)] = _ExtensionReader.from_bytes(
215 astropy.io.fits.BinTableHDU, tail_data[:json_size]
216 )
217 # Make any empty dictionary to cache deserialized objects. Keys are
218 # the zero-indexed row in the JSON table.
219 self._deserialized_pointer_cache: dict[int, Any] = {}
221 @classmethod
222 @contextmanager
223 def open(
224 cls,
225 path: ResourcePathExpression | IO[bytes],
226 *,
227 page_size: int = DEFAULT_PAGE_SIZE,
228 partial: bool = False,
229 ) -> Iterator[Self]:
230 """Create an output archive that writes to the given file.
232 Parameters
233 ----------
234 path
235 File to read; convertible to `lsst.resources.ResourcePath`,
236 or a seekable binary stream containing the file's content.
237 For stream input ``page_size`` and ``partial`` are ignored:
238 the data is already in memory and needs no paging.
239 page_size
240 Size of the fsspec read block for partial (remote) reads, in
241 bytes; a multiple of the FITS block size (2880) is recommended.
242 Defaults to `DEFAULT_PAGE_SIZE`; see it for the tuning tradeoff.
243 partial
244 Whether we will be reading only some of the archive, or if memory
245 pressure forces us to read it only a little at a time. If `False`
246 (default), the entire raw file may be read into memory up front.
248 Returns
249 -------
250 `contextlib.AbstractContextManager` [`FitsInputArchive`]
251 A context manager that returns a `FitsInputArchive` when entered.
252 """
253 if _is_binary_stream(path):
254 yield cls(path)
255 return
256 path = ResourcePath(path)
257 stream: IO[bytes]
258 if not partial:
259 stream = io.BytesIO(path.read())
260 yield cls(stream)
261 else:
262 fs: fsspec.AbstractFileSystem
263 fs, fp = path.to_fsspec()
264 # Cap cached blocks from the byte budget so memory stays bounded as
265 # the block size is retuned; keep at least two so the shared
266 # header/index block survives between a file's components.
267 maxblocks = max(2, READ_CACHE_MAX_BYTES // page_size)
268 with fs.open(
269 fp,
270 block_size=page_size,
271 cache_type=_READ_CACHE_TYPE,
272 cache_options={"maxblocks": maxblocks},
273 ) as stream:
274 yield cls(stream)
276 @property
277 def info(self) -> ArchiveInfo:
278 """Schema/format info read from the primary header on open
279 (`.serialization.ArchiveInfo`).
280 """
281 if self._info is None: 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true
282 raise ArchiveReadError("This is not an lsst.images FITS archive (no DATAMODL card).")
283 return self._info
285 def get_tree[T: ArchiveTree](self, model_type: type[T]) -> T:
286 """Read the JSON tree from the archive.
288 Parameters
289 ----------
290 model_type
291 A Pydantic model type to use to validate the JSON.
293 Returns
294 -------
295 T
296 The validated Pydantic model.
297 """
298 json_bytes = self._readers[ExtensionKey(JSON_EXTNAME)].data[0][JSON_COLUMN].tobytes()
299 return model_type.model_validate_json(json_bytes)
301 def deserialize_pointer[U: ArchiveTree, V](
302 self,
303 pointer: PointerModel,
304 model_type: type[U],
305 deserializer: Callable[[U, InputArchive[PointerModel]], V],
306 ) -> V:
307 # Docstring inherited.
308 if (cached := self._deserialized_pointer_cache.get(pointer.row)) is not None:
309 return cached
310 if not isinstance(pointer.column.data, ArrayReferenceModel):
311 raise ArchiveReadError(f"Invalid pointer with inline array:\n{pointer.model_dump_json(indent=2)}")
312 _, reader = self._get_source_reader(pointer.column.data.source, is_table=True)
313 try:
314 json_bytes = reader.data[pointer.row][JSON_COLUMN].tobytes()
315 except Exception as err:
316 raise InvalidFitsArchiveError(
317 f"Failed to access the table cell referenced by {pointer.model_dump_json()}."
318 ) from err
319 result = deserializer(model_type.model_validate_json(json_bytes), self)
320 self._deserialized_pointer_cache[pointer.row] = result
321 return result
323 def get_frame_set(self, ref: PointerModel) -> FrameSet:
324 try:
325 result = self._deserialized_pointer_cache[ref.row]
326 except KeyError:
327 raise AssertionError(
328 f"Frame set at {ref.model_dump_json(indent=2)} must be deserialized "
329 "before any dependent transform can be."
330 ) from None
331 if not isinstance(result, FrameSet):
332 raise InvalidFitsArchiveError(f"Expected a FrameSet instance at {ref.model_dump_json(indent=2)}.")
333 return result
335 def get_array(
336 self,
337 model: ArrayReferenceModel | InlineArrayModel,
338 *,
339 slices: tuple[slice, ...] | EllipsisType = ...,
340 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
341 ) -> np.ndarray:
342 if not isinstance(model, ArrayReferenceModel): 342 ↛ 343line 342 didn't jump to line 343 because the condition on line 342 was never true
343 raise ArchiveReadError("Inline array found where a reference array was expected.")
344 key, reader = self._get_source_reader(model.source, is_table=False)
345 if slices is not ...:
346 array = reader.section[slices]
347 else:
348 array = reader.data
349 if key not in self._opaque_metadata.headers: 349 ↛ 353line 349 didn't jump to line 353 because the condition on line 349 was always true
350 opaque_header = reader.header.copy(strip=True)
351 strip_header(opaque_header)
352 self._opaque_metadata.add_header(opaque_header, key=key)
353 return array
355 def get_table(
356 self,
357 model: TableModel,
358 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
359 ) -> astropy.table.Table:
360 # Docstring inherited.
361 array = self.get_structured_array(model, strip_header)
362 table = astropy.table.Table(array)
363 for c in model.columns:
364 c.update_table(table)
365 return table
367 def get_structured_array(
368 self,
369 model: TableModel,
370 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
371 ) -> np.ndarray:
372 # Docstring inherited.
373 if not isinstance(model.columns[0].data, ArrayReferenceModel):
374 raise ArchiveReadError("Inline array found where a reference array was expected.")
375 # All columns should have the same data.source; just use the first.
376 key, reader = self._get_source_reader(model.columns[0].data.source, is_table=True)
377 if key not in self._opaque_metadata.headers:
378 opaque_header = reader.header.copy(strip=True)
379 strip_header(opaque_header)
380 self._opaque_metadata.add_header(opaque_header, key=key)
381 return reader.hdu.data
383 def get_opaque_metadata(self) -> FitsOpaqueMetadata:
384 # Docstring inherited.
385 return self._opaque_metadata
387 def _get_source_reader(self, source: str | int, is_table: bool) -> tuple[ExtensionKey, _ExtensionReader]:
388 """Get a reader for the extension referenced by a serialiation model's
389 ``source`` field.
391 Parameters
392 ----------
393 source
394 A ``source`` field of the form ``fits:${hdu}`` or
395 ``fits:${hdu}[${col}]``.
396 is_table
397 Whether the source should be for a table HDU.
399 Returns
400 -------
401 key
402 Identifier pair for the HDU (EXTNAME, EXTVER).
403 reader
404 A reader object for the extension.
405 """
406 if not isinstance(source, str): 406 ↛ 407line 406 didn't jump to line 407 because the condition on line 406 was never true
407 raise InvalidFitsArchiveError(f"Reference with source={source!r} is not a string.")
408 if not source.startswith("fits:"): 408 ↛ 409line 408 didn't jump to line 409 because the condition on line 408 was never true
409 raise InvalidFitsArchiveError(f"Reference with source={source!r} does not start with 'fits:'.")
410 key = ExtensionKey.from_str(source)
411 try:
412 reader = self._readers[key]
413 except KeyError:
414 raise InvalidFitsArchiveError(f"Unrecognized source value {key}.") from None
415 if is_table and not reader.is_table: 415 ↛ 416line 415 didn't jump to line 416 because the condition on line 415 was never true
416 raise InvalidFitsArchiveError(
417 f"Extension with source={key} was expected to be be a binary table, not an image."
418 )
419 elif not is_table and reader.is_table: 419 ↛ 420line 419 didn't jump to line 420 because the condition on line 419 was never true
420 raise InvalidFitsArchiveError(
421 f"Extension with source={key} was expected to be be an image, not a binary table."
422 )
423 return key, reader
426class _ExtensionReader:
427 """A lazy-load reader for a single extension HDU.
429 Parameters
430 ----------
431 hdu_cls
432 The type of the astropy HDU instance to construct.
433 stream
434 The file-like object to read from.
435 """
437 def __init__(self, hdu_cls: type[ExtensionHDU], stream: IO[bytes]) -> None:
438 self._hdu_cls = hdu_cls
439 self._stream = stream
441 @classmethod
442 def from_index_row(cls, index_row: np.void, stream: IO[bytes]) -> _ExtensionReader:
443 """Construct from a row of the binary table index HDU.
445 Parameters
446 ----------
447 index_row
448 A record array row from the index HDU.
449 stream
450 The file-like object being used to read the full FITS file.
452 Returns
453 -------
454 reader
455 A reader object for the extension.
456 """
457 match index_row["XTENSION"].strip():
458 case "IMAGE":
459 hdu_cls = astropy.io.fits.ImageHDU
460 case "BINTABLE": 460 ↛ 465line 460 didn't jump to line 465 because the pattern on line 460 always matched
461 if index_row["ZIMAGE"]:
462 hdu_cls = astropy.io.fits.CompImageHDU
463 else:
464 hdu_cls = astropy.io.fits.BinTableHDU
465 case other:
466 raise AssertionError(f"Unsupported HDU type {other!r}.")
467 return _ExtensionReader(
468 hdu_cls,
469 _RangeStreamProxy(
470 stream,
471 start=int(index_row["HDRADDR"]),
472 ),
473 )
475 @classmethod
476 def from_bytes(cls, hdu_cls: type[ExtensionHDU], data: bytes) -> _ExtensionReader:
477 """Construct from already-read `bytes`.
479 Parameters
480 ----------
481 hdu_cls
482 The HDU type to instantiate.
483 data
484 Raw data for the HDU.
486 Returns
487 -------
488 reader
489 A reader object for extension.
490 """
491 return _ExtensionReader(hdu_cls, io.BytesIO(data))
493 @property
494 def is_table(self) -> bool:
495 """Whether this is logically a table HDU.
497 This is `False` for compressed image HDUs, even though they are
498 represented in FITS as a binary table.
499 """
500 return issubclass(self._hdu_cls, astropy.io.fits.BinTableHDU) and not issubclass(
501 self._hdu_cls, astropy.io.fits.CompImageHDU
502 )
504 @cached_property
505 def hdu(self) -> ExtensionHDU:
506 """The Astropy HDU object."""
507 self._stream.seek(0)
508 if self._hdu_cls is astropy.io.fits.CompImageHDU:
509 # CompImageHDU.readfrom doesn't work; we need to make a minimal
510 # example and report it upstream. Happily this workaround does
511 # work.
512 bintable_hdu = astropy.io.fits.BinTableHDU.readfrom(self._stream, memmap=False, cache=False)
513 return self._hdu_cls(bintable=bintable_hdu)
514 else:
515 return self._hdu_cls.readfrom(self._stream, memmap=False, cache=False, uint=True)
517 @property
518 def header(self) -> astropy.io.fits.Header:
519 """The header of the HDU."""
520 return self.hdu.header
522 @property
523 def data(self) -> np.ndarray:
524 """The data for the HDU."""
525 return self.hdu.data
527 @property
528 def section(self) -> astropy.io.fits.Section | astropy.io.fits.CompImageSection:
529 """An Astropy expression object that reads a subset of the data when
530 sliced.
531 """
532 return self.hdu.section
535class _RangeStreamProxy(IO[bytes]):
536 """A readable IO proxy object that makes the beginning of the file appear
537 at a custom position.
539 Parameters
540 ----------
541 base
542 Underlying readable, seekable buffer to proxy.
543 start
544 Offset into the base stream that will be considered the start of the
545 proxy stream.
547 Notes
548 -----
549 This class exists because Astropy doesn't seem to provide a way to read a
550 single HDU that starts at the current seek position of a file-like object.
551 It does provide ``readfrom`` methods on its HDU objects that take a
552 file-like object, but these assume (possibly unintentionally; it only
553 happens when Astropy is trying to see whether the file was opened for
554 appending) that ``seek(0)`` will set the file-like object to the start of
555 the HDU.
556 """
558 def __init__(self, base: IO[bytes], start: int) -> None:
559 self._base = base
560 self._start = start
562 @property
563 def mode(self) -> str:
564 return "rb"
566 def __enter__(self) -> Self:
567 raise AssertionError("This proxy should not be used as a context manager.")
569 def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
570 raise AssertionError("This proxy should not be used as a context manager.")
572 def __iter__(self) -> Iterator[bytes]:
573 return self._base.__iter__()
575 def __next__(self) -> bytes:
576 return self._base.__next__()
578 def close(self) -> None:
579 raise AssertionError("This proxy should not ever be closed.")
581 @property
582 def closed(self) -> bool:
583 return False
585 def fileno(self) -> int:
586 raise OSError()
588 def flush(self) -> None:
589 pass
591 def isatty(self) -> bool:
592 return False
594 def read(self, n: int = -1, /) -> bytes:
595 result = self._base.read(n)
596 return result
598 def readable(self) -> bool:
599 return True
601 def readline(self, limit: int = -1, /) -> bytes:
602 return self._base.readline(limit)
604 def readlines(self, hint: int = -1, /) -> list[bytes]:
605 return self._base.readlines(hint)
607 def seek(self, offset: int, whence: int = 0) -> int:
608 match whence:
609 case os.SEEK_SET:
610 return self._base.seek(offset + self._start, os.SEEK_SET) - self._start
611 case os.SEEK_CUR: 611 ↛ 612line 611 didn't jump to line 612 because the pattern on line 611 never matched
612 return self._base.seek(offset, os.SEEK_CUR) - self._start
613 case os.SEEK_END: 613 ↛ 615line 613 didn't jump to line 615 because the pattern on line 613 always matched
614 return self._base.seek(offset, os.SEEK_END) - self._start
615 raise TypeError(f"Invalid value for 'whence': {whence}.")
617 def seekable(self) -> bool:
618 return True
620 def tell(self) -> int:
621 return self._base.tell() - self._start
623 def truncate(self, size: int | None = None, /) -> int:
624 raise OSError()
626 def writable(self) -> bool:
627 return False
629 def write(self, arg: Any, /) -> int:
630 raise OSError()
632 def writelines(self, arg: Any, /) -> None:
633 raise OSError()