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