Coverage for python / lsst / images / fits / _input_archive.py: 30%
202 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-02 07:17 -0800
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 "read",
18)
20import io
21import os
22from collections.abc import Callable, Iterator
23from contextlib import contextmanager
24from functools import cached_property
25from types import EllipsisType
26from typing import IO, Any, Self
28import astropy.io.fits
29import astropy.table
30import fsspec
31import numpy as np
32import pydantic
34from lsst.resources import ResourcePath, ResourcePathExpression
36from .._transforms import FrameSet
37from ..serialization import (
38 ArchiveTree,
39 ArrayReferenceModel,
40 InputArchive,
41 TableCellReferenceModel,
42 TableModel,
43 no_header_updates,
44)
45from ._common import ExtensionHDU, ExtensionKey, FitsOpaqueMetadata, InvalidFitsArchiveError
48def read[T: Any](
49 cls: type[T],
50 path: ResourcePathExpression,
51 *,
52 page_size: int = 2880 * 50,
53 partial: bool | None = None,
54 **kwargs: Any,
55) -> T:
56 """Read an object from a FITS file.
58 Parameters
59 ----------
60 path
61 File to read; convertible to `lsst.resources.ResourcePath`.
62 page_size
63 Minimum number of bytes to read at at once. Making this a multiple
64 of the FITS block size (2880) is recommended.
65 partial
66 Whether we will be reading only some of the archive, or if memory
67 pressure forces us to read it only a little at a time. If `False`,
68 the entire raw file may be read into memory up front. Defaults to
69 `True` if any extra ``**kwargs`` are passed with values other than
70 `None`, since those usually indicate that only some of the original
71 object will be loaded.
72 **kwargs
73 Extra keyword arguments passed to ``cls.deserialize``.
75 Returns
76 -------
77 object
78 The loaded object.
80 Notes
81 -----
82 Supported types must implement ``deserialize`` and
83 ``_get_archive_tree_type`` (see `.Image` for an example).
84 """
85 if partial is None:
86 partial = any(v is not None for v in kwargs.values())
87 with FitsInputArchive.open(path, page_size=page_size, partial=partial) as archive:
88 tree = archive.get_tree(cls._get_archive_tree_type(TableCellReferenceModel))
89 obj = cls.deserialize(tree, archive, **kwargs)
90 obj._opaque_metadata = archive.get_opaque_metadata()
91 return obj
94class FitsInputArchive(InputArchive[TableCellReferenceModel]):
95 """An implementation of the `.serialization.InputArchive` interface that
96 reads from FITS files.
98 Instances of this class should only be constructed via the `open`
99 context manager.
100 """
102 def __init__(self, stream: IO[bytes]):
103 self._primary_hdu = astropy.io.fits.PrimaryHDU.readfrom(stream)
104 # TODO: read and strip subformat declaration and version, once we start
105 # writing those.
106 #
107 # TODO: do some basic checks that the file format conforms to our
108 # expectations (e.g. primary HDU should have no data).
109 #
110 # Read and strip the addresses and sizes from the headers. We don't
111 # actually need the index address because we always want to read the
112 # JSON HDU, too, and the index HDU is always the next one (but this
113 # could change in the future).
114 json_address: int = self._primary_hdu.header.pop("JSONADDR")
115 json_size: int = self._primary_hdu.header.pop("JSONSIZE")
116 del self._primary_hdu.header["INDXADDR"]
117 index_size: int = self._primary_hdu.header.pop("INDXSIZE")
118 # Save the remaining primary header keys so we can propagate them on
119 # rewrite.
120 self._opaque_metadata = FitsOpaqueMetadata()
121 self._opaque_metadata.headers[ExtensionKey()] = self._primary_hdu.header.copy(strip=True)
122 # Read the JSON and index HDUs from the end.
123 stream.seek(json_address)
124 tail_data = stream.read(json_size + index_size)
125 index_hdu = astropy.io.fits.BinTableHDU.fromstring(tail_data[json_size:])
126 # Initialize lazy readers for all of the regular HDUs and the JSON HDU.
127 self._readers = {
128 ExtensionKey.from_index_row(row): _ExtensionReader.from_index_row(row, stream)
129 for row in index_hdu.data
130 }
131 self._readers[ExtensionKey("JSON")] = _ExtensionReader.from_bytes(
132 astropy.io.fits.BinTableHDU, tail_data[:json_size]
133 )
134 # Make any empty dictionary to cache deserialized objects.
135 self._deserialized_pointer_cache: dict[TableCellReferenceModel, Any] = {}
137 @classmethod
138 @contextmanager
139 def open(
140 cls,
141 path: ResourcePathExpression,
142 *,
143 page_size: int = 2880 * 50,
144 partial: bool = False,
145 ) -> Iterator[Self]:
146 """Create an output archive that writes to the given file.
148 Parameters
149 ----------
150 path
151 File to read; convertible to `lsst.resources.ResourcePath`.
152 page_size
153 Minimum number of bytes to read at at once. Making this a multiple
154 of the FITS block size (2880) is recommended.
155 partial
156 Whether we will be reading only some of the archive, or if memory
157 pressure forces us to read it only a little at a time. If `False`
158 (default), the entire raw file may be read into memory up front.
160 Returns
161 -------
162 `contextlib.AbstractContextManager` [`FitsInputArchive`]
163 A context manager that returns a `FitsInputArchive` when entered.
164 """
165 path = ResourcePath(path)
166 stream: IO[bytes]
167 if not partial:
168 stream = io.BytesIO(path.read())
169 yield cls(stream)
170 else:
171 fs: fsspec.AbstractFileSystem
172 fs, fp = path.to_fsspec()
173 with fs.open(fp, block_size=page_size) as stream:
174 yield cls(stream)
176 def get_tree[T: ArchiveTree](self, model_type: type[T]) -> T:
177 """Read the JSON tree from the archive.
179 Parameters
180 ----------
181 model_type
182 A Pydantic model type to use to validate the JSON.
184 Returns
185 -------
186 T
187 The validated Pydantic model.
188 """
189 json_bytes = self._readers[ExtensionKey("JSON")].data[0]["JSON"].tobytes()
190 return model_type.model_validate_json(json_bytes)
192 def deserialize_pointer[U: pydantic.BaseModel, V](
193 self,
194 pointer: TableCellReferenceModel,
195 model_type: type[U],
196 deserializer: Callable[[U, InputArchive[TableCellReferenceModel]], V],
197 ) -> V:
198 # Docstring inherited.
199 if (cached := self._deserialized_pointer_cache.get(pointer)) is not None:
200 return cached
201 _, reader = self._get_source_reader(pointer)
202 try:
203 json_bytes = reader.data[pointer.row][pointer.column].tobytes()
204 except Exception as err:
205 raise InvalidFitsArchiveError(
206 f"Failed to access the table cell referenced by {pointer.model_dump_json()}."
207 ) from err
208 result = deserializer(model_type.model_validate_json(json_bytes), self)
209 self._deserialized_pointer_cache[pointer] = result
210 return result
212 def get_frame_set(self, ref: TableCellReferenceModel) -> FrameSet:
213 try:
214 result = self._deserialized_pointer_cache[ref]
215 except KeyError:
216 raise AssertionError(
217 f"Frame set at {ref.model_dump_json(indent=2)} must be deserialized "
218 "before any dependent transform can be."
219 ) from None
220 if not isinstance(result, FrameSet):
221 raise InvalidFitsArchiveError(f"Expected a FrameSet instance at {ref.model_dump_json(indent=2)}.")
222 return result
224 def get_array(
225 self,
226 ref: ArrayReferenceModel,
227 *,
228 slices: tuple[slice, ...] | EllipsisType = ...,
229 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
230 ) -> np.ndarray:
231 key, reader = self._get_source_reader(ref)
232 if slices is not ...:
233 array = reader.section[slices]
234 else:
235 array = reader.data
236 if key not in self._opaque_metadata.headers:
237 opaque_header = reader.header.copy(strip=True)
238 strip_header(opaque_header)
239 self._opaque_metadata.add_header(opaque_header, key=key)
240 return array
242 def get_table(
243 self,
244 ref: TableModel,
245 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
246 ) -> astropy.table.Table:
247 # Docstring inherited.
248 array = self.get_structured_array(ref, strip_header)
249 table = astropy.table.Table(array)
250 for c in ref.columns:
251 c.update_table(table)
252 return table
254 def get_structured_array(
255 self,
256 ref: TableModel,
257 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
258 ) -> np.ndarray:
259 # Docstring inherited.
260 key, reader = self._get_source_reader(ref)
261 if key not in self._opaque_metadata.headers:
262 opaque_header = reader.header.copy(strip=True)
263 strip_header(opaque_header)
264 self._opaque_metadata.add_header(opaque_header, key=key)
265 return reader.hdu.data
267 def get_opaque_metadata(self) -> FitsOpaqueMetadata:
268 # Docstring inherited.
269 return self._opaque_metadata
271 def _get_source_reader(
272 self, ref: ArrayReferenceModel | TableCellReferenceModel | TableModel
273 ) -> tuple[ExtensionKey, _ExtensionReader]:
274 """Get a reader for the extension referenced by a model.
276 Parameters
277 ----------
278 ref
279 A Pydantic model with a 'source' field that references a FITS HDU.
281 Returns
282 -------
283 key
284 Identifier pair for the HDU (EXTNAME, EXTVER).
285 reader
286 A reader object for the extension.
287 """
288 if not isinstance(ref.source, str):
289 raise InvalidFitsArchiveError(f"Reference with source={ref.source!r} is not a string.")
290 if not ref.source.startswith("fits:"):
291 raise InvalidFitsArchiveError(
292 f"Reference with source={ref.source!r} does not start with 'fits:'."
293 )
294 key = ExtensionKey.from_str(ref.source)
295 try:
296 reader = self._readers[key]
297 except KeyError:
298 raise InvalidFitsArchiveError(f"Unrecognized source value {key}.") from None
299 if ref.source_is_table and not reader.is_table:
300 raise InvalidFitsArchiveError(
301 f"Extension with source={key} was expected to be be a binary table, not an image."
302 )
303 elif not ref.source_is_table and reader.is_table:
304 raise InvalidFitsArchiveError(
305 f"Extension with source={key} was expected to be be an image, not a binary table."
306 )
307 return key, reader
310class _ExtensionReader:
311 """A lazy-load reader for a single extension HDU.
313 Parameters
314 ----------
315 hdu_cls
316 The type of the astropy HDU instance to construct.
317 stream
318 The file-like object to read from.
319 """
321 def __init__(self, hdu_cls: type[ExtensionHDU], stream: IO[bytes]):
322 self._hdu_cls = hdu_cls
323 self._stream = stream
325 @classmethod
326 def from_index_row(cls, index_row: np.void, stream: IO[bytes]) -> _ExtensionReader:
327 """Construct from a row of the binary table index HDU.
329 Parameters
330 ----------
331 index_row
332 A record array row from the index HDU.
333 stream
334 The file-like object being used to read the full FITS file.
336 Returns
337 -------
338 reader
339 A reader object for the extension.
340 """
341 match index_row["XTENSION"].strip():
342 case "IMAGE":
343 hdu_cls = astropy.io.fits.ImageHDU
344 case "BINTABLE":
345 if index_row["ZIMAGE"]:
346 hdu_cls = astropy.io.fits.CompImageHDU
347 else:
348 hdu_cls = astropy.io.fits.BinTableHDU
349 case other:
350 raise AssertionError(f"Unsupported HDU type {other!r}.")
351 return _ExtensionReader(
352 hdu_cls,
353 _RangeStreamProxy(
354 stream,
355 start=int(index_row["HDRADDR"]),
356 ),
357 )
359 @classmethod
360 def from_bytes(cls, hdu_cls: type[ExtensionHDU], data: bytes) -> _ExtensionReader:
361 """Construct from already-read `bytes`
363 Parameters
364 ----------
365 hdu_cls
366 The HDU type to instantiate.
367 data
368 Raw data for the HDU.
370 Returns
371 -------
372 reader
373 A reader object for extension.
374 """
375 return _ExtensionReader(hdu_cls, io.BytesIO(data))
377 @property
378 def is_table(self) -> bool:
379 """Whether this is logically a table HDU.
381 This is `False` for compressed image HDUs, even though they are
382 represented in FITS as a binary table.
383 """
384 return issubclass(self._hdu_cls, astropy.io.fits.BinTableHDU) and not issubclass(
385 self._hdu_cls, astropy.io.fits.CompImageHDU
386 )
388 @cached_property
389 def hdu(self) -> ExtensionHDU:
390 """The Astropy HDU object."""
391 self._stream.seek(0)
392 if self._hdu_cls is astropy.io.fits.CompImageHDU:
393 # CompImageHDU.readfrom doesn't work; we need to make a minimal
394 # example and report it upstream. Happily this workaround does
395 # work.
396 bintable_hdu = astropy.io.fits.BinTableHDU.readfrom(self._stream, memmap=False, cache=False)
397 return self._hdu_cls(bintable=bintable_hdu)
398 else:
399 return self._hdu_cls.readfrom(self._stream, memmap=False, cache=False)
401 @property
402 def header(self) -> astropy.io.fits.Header:
403 """The header of the HDU."""
404 return self.hdu.header
406 @property
407 def data(self) -> np.ndarray:
408 """The data for the HDU."""
409 return self.hdu.data
411 @property
412 def section(self) -> astropy.io.fits.Section | astropy.io.fits.CompImageSection:
413 """An Astropy expression object that reads a subset of the data when
414 sliced.
415 """
416 return self.hdu.section
419class _RangeStreamProxy(IO[bytes]):
420 """A readable IO proxy object that makes the beginning of the file appear
421 at a custom position.
423 Parameters
424 ----------
425 base
426 Underlying readable, seekable buffer to proxy.
427 start
428 Offset into the base stream that will be considered the start of the
429 proxy stream.
431 Notes
432 -----
433 This class exists because Astropy doesn't seem to provide a way to read a
434 single HDU that starts at the current seek position of a file-like object.
435 It does provide ``readfrom`` methods on its HDU objects that take a
436 file-like object, but these assume (possibly unintentionally; it only
437 happens when Astropy is trying to see whether the file was opened for
438 appending) that ``seek(0)`` will set the file-like object to the start of
439 the HDU.
440 """
442 def __init__(self, base: IO[bytes], start: int):
443 self._base = base
444 self._start = start
446 @property
447 def mode(self) -> str:
448 return "rb"
450 def __enter__(self) -> Self:
451 raise AssertionError("This proxy should not be used as a context manager.")
453 def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
454 raise AssertionError("This proxy should not be used as a context manager.")
456 def __iter__(self) -> Iterator[bytes]:
457 return self._base.__iter__()
459 def __next__(self) -> bytes:
460 return self._base.__next__()
462 def close(self) -> None:
463 raise AssertionError("This proxy should not ever be closed.")
465 @property
466 def closed(self) -> bool:
467 return False
469 def fileno(self) -> int:
470 raise OSError()
472 def flush(self) -> None:
473 pass
475 def isatty(self) -> bool:
476 return False
478 def read(self, n: int = -1, /) -> bytes:
479 result = self._base.read(n)
480 return result
482 def readable(self) -> bool:
483 return True
485 def readline(self, limit: int = -1, /) -> bytes:
486 return self._base.readline(limit)
488 def readlines(self, hint: int = -1, /) -> list[bytes]:
489 return self._base.readlines(hint)
491 def seek(self, offset: int, whence: int = 0) -> int:
492 match whence:
493 case os.SEEK_SET:
494 return self._base.seek(offset + self._start, os.SEEK_SET) - self._start
495 case os.SEEK_CUR:
496 return self._base.seek(offset, os.SEEK_CUR) - self._start
497 case os.SEEK_END:
498 return self._base.seek(offset, os.SEEK_END) - self._start
499 raise TypeError(f"Invalid value for 'whence': {whence}.")
501 def seekable(self) -> bool:
502 return True
504 def tell(self) -> int:
505 return self._base.tell() - self._start
507 def truncate(self, size: int | None = None, /) -> int:
508 raise OSError()
510 def writable(self) -> bool:
511 return False
513 def write(self, arg: Any, /) -> int:
514 raise OSError()
516 def writelines(self, arg: Any, /) -> None:
517 raise OSError()