Coverage for python/lsst/images/fits/_input_archive.py: 30%
223 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-06 01:43 -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 "FitsInputArchive",
16 "FitsOpaqueMetadata",
17)
19import io
20import os
21from collections.abc import Callable, Iterator
22from contextlib import contextmanager
23from functools import cached_property
24from types import EllipsisType
25from typing import IO, Any, Self
27import astropy.io.fits
28import astropy.table
29import fsspec
30import numpy as np
32from lsst.resources import ResourcePath, ResourcePathExpression
34from .._transforms import FrameSet
35from ..serialization import (
36 ArchiveInfo,
37 ArchiveReadError,
38 ArchiveTree,
39 ArrayReferenceModel,
40 InlineArrayModel,
41 InputArchive,
42 TableModel,
43 no_header_updates,
44 parameterize_tree,
45)
46from ..serialization._common import _check_format_version
47from ._common import (
48 JSON_COLUMN,
49 JSON_EXTNAME,
50 ExtensionHDU,
51 ExtensionKey,
52 FitsOpaqueMetadata,
53 InvalidFitsArchiveError,
54 PointerModel,
55)
57_FITS_FORMAT_VERSION = 1
58"""Container layout version this release of `FitsInputArchive` understands."""
61class FitsInputArchive(InputArchive[PointerModel]):
62 """An implementation of the `.serialization.InputArchive` interface that
63 reads from FITS files.
65 Instances of this class should only be constructed via the `open`
66 context manager.
67 """
69 @classmethod
70 def get_basic_info(cls, path: ResourcePathExpression) -> ArchiveInfo:
71 """Read ``DATAMODL`` (schema URL) and ``FMTVER`` (container version)
72 from the primary header.
74 Every FITS file written by this package records the schema URL in
75 the ``DATAMODL`` card, so the schema can be identified without
76 reading the (potentially large) JSON tree HDU.
77 """
78 with ResourcePath(path).open("rb") as stream:
79 primary = astropy.io.fits.PrimaryHDU.readfrom(stream)
80 header = primary.header
81 format_version = int(header.get("FMTVER", 1))
82 schema_url = header.get("DATAMODL")
83 if not schema_url:
84 raise ArchiveReadError(f"{path!r} is not an lsst.images FITS archive (no DATAMODL card).")
85 return ArchiveInfo.from_schema_url(schema_url, format_version=format_version)
87 @classmethod
88 @contextmanager
89 def open_tree(
90 cls,
91 path: ResourcePathExpression,
92 tree_cls: type[ArchiveTree],
93 *,
94 partial: bool = True,
95 **backend_kwargs: Any,
96 ) -> Iterator[tuple[Self, ArchiveTree]]:
97 """Open the FITS file and yield ``(archive, tree)``.
99 Honors the ``page_size`` and ``partial`` open options.
100 """
101 page_size = backend_kwargs.pop("page_size", 2880 * 50)
102 parameterized = parameterize_tree(tree_cls, PointerModel)
103 with cls.open(path, page_size=page_size, partial=partial) as archive:
104 tree = archive.get_tree(parameterized)
105 yield archive, tree
107 def __init__(self, stream: IO[bytes]):
108 self._primary_hdu = astropy.io.fits.PrimaryHDU.readfrom(stream)
109 on_disk_fmtver: int = self._primary_hdu.header.pop("FMTVER", 1)
110 # DATAMODL is informational only on read; the JSON tree's
111 # schema_version / min_read_version drive data-model checks.
112 self._primary_hdu.header.pop("DATAMODL", None)
113 _check_format_version("fits", on_disk_fmtver, _FITS_FORMAT_VERSION)
114 # TODO: do some basic checks that the file format conforms to our
115 # expectations (e.g. primary HDU should have no data).
116 #
117 # Read and strip the addresses and sizes from the headers. We don't
118 # actually need the index address because we always want to read the
119 # JSON HDU, too, and the index HDU is always the next one (but this
120 # could change in the future).
121 json_address: int = self._primary_hdu.header.pop("JSONADDR")
122 json_size: int = self._primary_hdu.header.pop("JSONSIZE")
123 del self._primary_hdu.header["INDXADDR"]
124 index_size: int = self._primary_hdu.header.pop("INDXSIZE")
125 # Save the remaining primary header keys so we can propagate them on
126 # rewrite.
127 self._opaque_metadata = FitsOpaqueMetadata()
128 self._opaque_metadata.add_header(self._primary_hdu.header.copy(strip=True), name="", ver=1)
129 # Read the JSON and index HDUs from the end.
130 stream.seek(json_address)
131 tail_data = stream.read(json_size + index_size)
132 index_hdu = astropy.io.fits.BinTableHDU.fromstring(tail_data[json_size:])
133 # Initialize lazy readers for all of the regular HDUs and the JSON HDU.
134 self._readers = {
135 ExtensionKey.from_index_row(row): _ExtensionReader.from_index_row(row, stream)
136 for row in index_hdu.data
137 }
138 self._readers[ExtensionKey(JSON_COLUMN)] = _ExtensionReader.from_bytes(
139 astropy.io.fits.BinTableHDU, tail_data[:json_size]
140 )
141 # Make any empty dictionary to cache deserialized objects. Keys are
142 # the zero-indexed row in the JSON table.
143 self._deserialized_pointer_cache: dict[int, Any] = {}
145 @classmethod
146 @contextmanager
147 def open(
148 cls,
149 path: ResourcePathExpression,
150 *,
151 page_size: int = 2880 * 50,
152 partial: bool = False,
153 ) -> Iterator[Self]:
154 """Create an output archive that writes to the given file.
156 Parameters
157 ----------
158 path
159 File to read; convertible to `lsst.resources.ResourcePath`.
160 page_size
161 Minimum number of bytes to read at at once. Making this a multiple
162 of the FITS block size (2880) is recommended.
163 partial
164 Whether we will be reading only some of the archive, or if memory
165 pressure forces us to read it only a little at a time. If `False`
166 (default), the entire raw file may be read into memory up front.
168 Returns
169 -------
170 `contextlib.AbstractContextManager` [`FitsInputArchive`]
171 A context manager that returns a `FitsInputArchive` when entered.
172 """
173 path = ResourcePath(path)
174 stream: IO[bytes]
175 if not partial:
176 stream = io.BytesIO(path.read())
177 yield cls(stream)
178 else:
179 fs: fsspec.AbstractFileSystem
180 fs, fp = path.to_fsspec()
181 with fs.open(fp, block_size=page_size) as stream:
182 yield cls(stream)
184 def get_tree[T: ArchiveTree](self, model_type: type[T]) -> T:
185 """Read the JSON tree from the archive.
187 Parameters
188 ----------
189 model_type
190 A Pydantic model type to use to validate the JSON.
192 Returns
193 -------
194 T
195 The validated Pydantic model.
196 """
197 json_bytes = self._readers[ExtensionKey(JSON_EXTNAME)].data[0][JSON_COLUMN].tobytes()
198 return model_type.model_validate_json(json_bytes)
200 def deserialize_pointer[U: ArchiveTree, V](
201 self,
202 pointer: PointerModel,
203 model_type: type[U],
204 deserializer: Callable[[U, InputArchive[PointerModel]], V],
205 ) -> V:
206 # Docstring inherited.
207 if (cached := self._deserialized_pointer_cache.get(pointer.row)) is not None:
208 return cached
209 if not isinstance(pointer.column.data, ArrayReferenceModel):
210 raise ArchiveReadError(f"Invalid pointer with inline array:\n{pointer.model_dump_json(indent=2)}")
211 _, reader = self._get_source_reader(pointer.column.data.source, is_table=True)
212 try:
213 json_bytes = reader.data[pointer.row][JSON_COLUMN].tobytes()
214 except Exception as err:
215 raise InvalidFitsArchiveError(
216 f"Failed to access the table cell referenced by {pointer.model_dump_json()}."
217 ) from err
218 result = deserializer(model_type.model_validate_json(json_bytes), self)
219 self._deserialized_pointer_cache[pointer.row] = result
220 return result
222 def get_frame_set(self, ref: PointerModel) -> FrameSet:
223 try:
224 result = self._deserialized_pointer_cache[ref.row]
225 except KeyError:
226 raise AssertionError(
227 f"Frame set at {ref.model_dump_json(indent=2)} must be deserialized "
228 "before any dependent transform can be."
229 ) from None
230 if not isinstance(result, FrameSet):
231 raise InvalidFitsArchiveError(f"Expected a FrameSet instance at {ref.model_dump_json(indent=2)}.")
232 return result
234 def get_array(
235 self,
236 model: ArrayReferenceModel | InlineArrayModel,
237 *,
238 slices: tuple[slice, ...] | EllipsisType = ...,
239 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
240 ) -> np.ndarray:
241 if not isinstance(model, ArrayReferenceModel):
242 raise ArchiveReadError("Inline array found where a reference array was expected.")
243 key, reader = self._get_source_reader(model.source, is_table=False)
244 if slices is not ...:
245 array = reader.section[slices]
246 else:
247 array = reader.data
248 if key not in self._opaque_metadata.headers:
249 opaque_header = reader.header.copy(strip=True)
250 strip_header(opaque_header)
251 self._opaque_metadata.add_header(opaque_header, key=key)
252 return array
254 def get_table(
255 self,
256 model: TableModel,
257 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
258 ) -> astropy.table.Table:
259 # Docstring inherited.
260 array = self.get_structured_array(model, strip_header)
261 table = astropy.table.Table(array)
262 for c in model.columns:
263 c.update_table(table)
264 return table
266 def get_structured_array(
267 self,
268 model: TableModel,
269 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
270 ) -> np.ndarray:
271 # Docstring inherited.
272 if not isinstance(model.columns[0].data, ArrayReferenceModel):
273 raise ArchiveReadError("Inline array found where a reference array was expected.")
274 # All columns should have the same data.source; just use the first.
275 key, reader = self._get_source_reader(model.columns[0].data.source, is_table=True)
276 if key not in self._opaque_metadata.headers:
277 opaque_header = reader.header.copy(strip=True)
278 strip_header(opaque_header)
279 self._opaque_metadata.add_header(opaque_header, key=key)
280 return reader.hdu.data
282 def get_opaque_metadata(self) -> FitsOpaqueMetadata:
283 # Docstring inherited.
284 return self._opaque_metadata
286 def _get_source_reader(self, source: str | int, is_table: bool) -> tuple[ExtensionKey, _ExtensionReader]:
287 """Get a reader for the extension referenced by a serialiation model's
288 ``source`` field.
290 Parameters
291 ----------
292 source
293 A ``source`` field of the form ``fits:${hdu}`` or
294 ``fits:${hdu}[${col}]``.
295 is_table
296 Whether the source should be for a table HDU.
298 Returns
299 -------
300 key
301 Identifier pair for the HDU (EXTNAME, EXTVER).
302 reader
303 A reader object for the extension.
304 """
305 if not isinstance(source, str):
306 raise InvalidFitsArchiveError(f"Reference with source={source!r} is not a string.")
307 if not source.startswith("fits:"):
308 raise InvalidFitsArchiveError(f"Reference with source={source!r} does not start with 'fits:'.")
309 key = ExtensionKey.from_str(source)
310 try:
311 reader = self._readers[key]
312 except KeyError:
313 raise InvalidFitsArchiveError(f"Unrecognized source value {key}.") from None
314 if is_table and not reader.is_table:
315 raise InvalidFitsArchiveError(
316 f"Extension with source={key} was expected to be be a binary table, not an image."
317 )
318 elif not is_table and reader.is_table:
319 raise InvalidFitsArchiveError(
320 f"Extension with source={key} was expected to be be an image, not a binary table."
321 )
322 return key, reader
325class _ExtensionReader:
326 """A lazy-load reader for a single extension HDU.
328 Parameters
329 ----------
330 hdu_cls
331 The type of the astropy HDU instance to construct.
332 stream
333 The file-like object to read from.
334 """
336 def __init__(self, hdu_cls: type[ExtensionHDU], stream: IO[bytes]):
337 self._hdu_cls = hdu_cls
338 self._stream = stream
340 @classmethod
341 def from_index_row(cls, index_row: np.void, stream: IO[bytes]) -> _ExtensionReader:
342 """Construct from a row of the binary table index HDU.
344 Parameters
345 ----------
346 index_row
347 A record array row from the index HDU.
348 stream
349 The file-like object being used to read the full FITS file.
351 Returns
352 -------
353 reader
354 A reader object for the extension.
355 """
356 match index_row["XTENSION"].strip():
357 case "IMAGE":
358 hdu_cls = astropy.io.fits.ImageHDU
359 case "BINTABLE":
360 if index_row["ZIMAGE"]:
361 hdu_cls = astropy.io.fits.CompImageHDU
362 else:
363 hdu_cls = astropy.io.fits.BinTableHDU
364 case other:
365 raise AssertionError(f"Unsupported HDU type {other!r}.")
366 return _ExtensionReader(
367 hdu_cls,
368 _RangeStreamProxy(
369 stream,
370 start=int(index_row["HDRADDR"]),
371 ),
372 )
374 @classmethod
375 def from_bytes(cls, hdu_cls: type[ExtensionHDU], data: bytes) -> _ExtensionReader:
376 """Construct from already-read `bytes`
378 Parameters
379 ----------
380 hdu_cls
381 The HDU type to instantiate.
382 data
383 Raw data for the HDU.
385 Returns
386 -------
387 reader
388 A reader object for extension.
389 """
390 return _ExtensionReader(hdu_cls, io.BytesIO(data))
392 @property
393 def is_table(self) -> bool:
394 """Whether this is logically a table HDU.
396 This is `False` for compressed image HDUs, even though they are
397 represented in FITS as a binary table.
398 """
399 return issubclass(self._hdu_cls, astropy.io.fits.BinTableHDU) and not issubclass(
400 self._hdu_cls, astropy.io.fits.CompImageHDU
401 )
403 @cached_property
404 def hdu(self) -> ExtensionHDU:
405 """The Astropy HDU object."""
406 self._stream.seek(0)
407 if self._hdu_cls is astropy.io.fits.CompImageHDU:
408 # CompImageHDU.readfrom doesn't work; we need to make a minimal
409 # example and report it upstream. Happily this workaround does
410 # work.
411 bintable_hdu = astropy.io.fits.BinTableHDU.readfrom(self._stream, memmap=False, cache=False)
412 return self._hdu_cls(bintable=bintable_hdu)
413 else:
414 return self._hdu_cls.readfrom(self._stream, memmap=False, cache=False, uint=True)
416 @property
417 def header(self) -> astropy.io.fits.Header:
418 """The header of the HDU."""
419 return self.hdu.header
421 @property
422 def data(self) -> np.ndarray:
423 """The data for the HDU."""
424 return self.hdu.data
426 @property
427 def section(self) -> astropy.io.fits.Section | astropy.io.fits.CompImageSection:
428 """An Astropy expression object that reads a subset of the data when
429 sliced.
430 """
431 return self.hdu.section
434class _RangeStreamProxy(IO[bytes]):
435 """A readable IO proxy object that makes the beginning of the file appear
436 at a custom position.
438 Parameters
439 ----------
440 base
441 Underlying readable, seekable buffer to proxy.
442 start
443 Offset into the base stream that will be considered the start of the
444 proxy stream.
446 Notes
447 -----
448 This class exists because Astropy doesn't seem to provide a way to read a
449 single HDU that starts at the current seek position of a file-like object.
450 It does provide ``readfrom`` methods on its HDU objects that take a
451 file-like object, but these assume (possibly unintentionally; it only
452 happens when Astropy is trying to see whether the file was opened for
453 appending) that ``seek(0)`` will set the file-like object to the start of
454 the HDU.
455 """
457 def __init__(self, base: IO[bytes], start: int):
458 self._base = base
459 self._start = start
461 @property
462 def mode(self) -> str:
463 return "rb"
465 def __enter__(self) -> Self:
466 raise AssertionError("This proxy should not be used as a context manager.")
468 def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
469 raise AssertionError("This proxy should not be used as a context manager.")
471 def __iter__(self) -> Iterator[bytes]:
472 return self._base.__iter__()
474 def __next__(self) -> bytes:
475 return self._base.__next__()
477 def close(self) -> None:
478 raise AssertionError("This proxy should not ever be closed.")
480 @property
481 def closed(self) -> bool:
482 return False
484 def fileno(self) -> int:
485 raise OSError()
487 def flush(self) -> None:
488 pass
490 def isatty(self) -> bool:
491 return False
493 def read(self, n: int = -1, /) -> bytes:
494 result = self._base.read(n)
495 return result
497 def readable(self) -> bool:
498 return True
500 def readline(self, limit: int = -1, /) -> bytes:
501 return self._base.readline(limit)
503 def readlines(self, hint: int = -1, /) -> list[bytes]:
504 return self._base.readlines(hint)
506 def seek(self, offset: int, whence: int = 0) -> int:
507 match whence:
508 case os.SEEK_SET:
509 return self._base.seek(offset + self._start, os.SEEK_SET) - self._start
510 case os.SEEK_CUR:
511 return self._base.seek(offset, os.SEEK_CUR) - self._start
512 case os.SEEK_END:
513 return self._base.seek(offset, os.SEEK_END) - self._start
514 raise TypeError(f"Invalid value for 'whence': {whence}.")
516 def seekable(self) -> bool:
517 return True
519 def tell(self) -> int:
520 return self._base.tell() - self._start
522 def truncate(self, size: int | None = None, /) -> int:
523 raise OSError()
525 def writable(self) -> bool:
526 return False
528 def write(self, arg: Any, /) -> int:
529 raise OSError()
531 def writelines(self, arg: Any, /) -> None:
532 raise OSError()