Coverage for python / lsst / images / fits / formatters.py: 0%
117 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-27 01:47 -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__ = ("GenericFormatter", "ImageFormatter", "MaskedImageFormatter", "VisitImageFormatter")
16import enum
17import hashlib
18import json
19from typing import Any, ClassVar
21import astropy.io.fits
23from lsst.daf.butler import DatasetProvenance, FormatterV2
24from lsst.resources import ResourcePath
26from .._geom import Box
27from .._image import Image
28from .._mask import Mask
29from .._masked_image import MaskedImageSerializationModel
30from .._transforms import Projection, ProjectionSerializationModel
31from .._visit_image import VisitImageSerializationModel
32from ..serialization import ButlerInfo, TableCellReferenceModel
33from ._common import FitsCompressionOptions
34from ._input_archive import FitsInputArchive, read
35from ._output_archive import write
38class GenericFormatter(FormatterV2):
39 """The butler interface to FITS archive serialization.
41 Serialized types must meet all the requirements of the `read` and `write`
42 functions.
44 Notes
45 -----
46 This formatter just forwards all read parameters it receives as
47 ``**kwargs`` to `.read` and hence the ``deserialize`` method of the type it
48 is reading. This may or may not be appropriate.
50 This formatter must be subclassed to add component support.
52 The write parameter configuration for this formatter is designed to be
53 identical to that for the legacy FITS formatters defined in
54 `lsst.obs.base`.
56 Butler provenance is written to both FITS headers and the archive tree.
57 """
59 default_extension: ClassVar[str] = ".fits"
60 can_read_from_uri: ClassVar[bool] = True
61 supported_write_parameters: ClassVar[frozenset[str]] = frozenset({"recipe"})
63 butler_provenance: DatasetProvenance | None = None
65 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any:
66 pytype = self.dataset_ref.datasetType.storageClass.pytype
67 kwargs = self.file_descriptor.parameters or {}
68 return read(pytype, uri, **kwargs).deserialized
70 def write_local_file(self, in_memory_dataset: Any, uri: ResourcePath) -> None:
71 butler_info = ButlerInfo(
72 dataset=self.dataset_ref.to_simple(),
73 provenance=self.butler_provenance if self.butler_provenance is not None else DatasetProvenance(),
74 )
75 write(
76 in_memory_dataset,
77 uri.ospath,
78 update_header=self._update_header,
79 compression_options=self._get_compression_options(),
80 compression_seed=self._get_compression_seed(),
81 butler_info=butler_info,
82 )
84 def add_provenance(
85 self, in_memory_dataset: Any, /, *, provenance: DatasetProvenance | None = None
86 ) -> Any:
87 # Instead of attaching the provenance to the object we remember it on
88 # the formatter, since a Formatter instance is only used once.
89 self.butler_provenance = provenance
90 return in_memory_dataset
92 def _get_compression_seed(self) -> int:
93 # Set the seed based on data ID (all logic here duplicated from
94 # obs_base). We can't just use 'hash', since like 'set' that's not
95 # deterministic. And we can't rely on a DimensionPacker because those
96 # are only defined for certain combinations of dimensions. Doing an MD5
97 # of the JSON feels like overkill but I don't really see anything much
98 # simpler.
99 hash_bytes = hashlib.md5(
100 json.dumps(list(self.data_id.required_values)).encode(),
101 usedforsecurity=False,
102 ).digest()
103 # And it *really* feels like overkill when we squash that into the [1,
104 # 10000] range allowed by FITS.
105 return 1 + int.from_bytes(hash_bytes) % 9999
107 def _get_compression_options(self) -> dict[str, FitsCompressionOptions]:
108 recipe = self.write_parameters.get("recipe", "default")
109 try:
110 config = self.write_recipes[recipe]
111 except KeyError:
112 if recipe == "default":
113 # If there's no default recipe just use the software defaults.
114 return {}
115 raise RuntimeError(f"Invalid recipe for ImageFormatter: {recipe!r}.") from None
116 return {k: FitsCompressionOptions.model_validate(v) for k, v in config.items()}
118 def _update_header(self, header: astropy.io.fits.Header) -> None:
119 # Logic here largely lifted from lsst.obs.base.utils, which we
120 # can't use directly for dependency and maybe mapping-type
121 # (PropertyList vs. astropy) reasons. We assume we can always add
122 # long cards (astropy will CONTINUE them) but not comments
123 # (astropy will truncate and warn on long cards).
124 for key in list(header):
125 if key.startswith("LSST BUTLER"):
126 del header[key]
127 if self.butler_provenance is not None:
128 for key, value in self.butler_provenance.to_flat_dict(
129 self.dataset_ref, prefix="HIERARCH LSST BUTLER", sep=" ", simple_types=True, max_inputs=3_000
130 ).items():
131 header.set(key, value)
134class ComponentSentinel(enum.Enum):
135 """Special values returned by `ImageFormatter.read_component`."""
137 UNRECOGNIZED_COMPONENT = enum.auto()
138 """This formatter does not recognize the given component, but a subclass
139 might.
140 """
142 INVALID_COMPONENT_MODEL = enum.auto()
143 """This formatter recognizes the given component, but the expected
144 attribute of the top-level `..serialization.ArchiveTree` did not exist
145 or had the wrong type.
146 """
149class ImageFormatter(GenericFormatter):
150 """The specialized butler interface to FITS archive serialization of
151 image-like objects with ``projection`` and ``bbox`` components.
153 Notes
154 -----
155 This formatter works by assuming the `..serialization.ArchiveTree` for the
156 top-level object has a ``projection`` attribute (a
157 `..ProjectionSerializationModel`) and a ``bbox`` property (a `..Box`).
159 Subclasses can add support for additional components by overriding
160 `read_component`, delegating to `super`, and handling the cases where it
161 returns a `ComponentSentinel` instance.
162 """
164 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any:
165 pytype: Any = self.file_descriptor.storageClass.pytype
166 if component is None:
167 result = read(pytype, uri, bbox=self.pop_bbox_from_parameters()).deserialized
168 else:
169 with FitsInputArchive.open(uri, partial=True) as archive:
170 tree = archive.get_tree(pytype._get_archive_tree_type(TableCellReferenceModel))
171 result = self.read_component(component, tree, archive)
172 if result is ComponentSentinel.UNRECOGNIZED_COMPONENT:
173 raise NotImplementedError(
174 f"Unrecognized component {component!r} for {type(self).__name__}."
175 )
176 if result is ComponentSentinel.INVALID_COMPONENT_MODEL:
177 raise NotImplementedError(
178 f"Invalid serialization model for component {component!r} for {type(self).__name__}."
179 )
180 self.check_unhandled_parameters()
181 return result
183 def pop_bbox_from_parameters(self) -> Box | None:
184 parameters = self.file_descriptor.parameters or {}
185 return parameters.pop("bbox", None)
187 def check_unhandled_parameters(self) -> None:
188 if self.file_descriptor.parameters:
189 raise RuntimeError(f"Parameters {list(self.file_descriptor.parameters.keys())} not recognized.")
191 def read_component(
192 self,
193 component: str,
194 tree: Any,
195 archive: FitsInputArchive,
196 ) -> Any:
197 match component:
198 case "projection":
199 if isinstance(
200 serialized_projection := getattr(tree, "projection", None), ProjectionSerializationModel
201 ):
202 return Projection.deserialize(serialized_projection, archive)
203 else:
204 return ComponentSentinel.INVALID_COMPONENT_MODEL
205 case "bbox":
206 if isinstance(bbox := getattr(tree, "bbox", None), Box):
207 return bbox
208 else:
209 return ComponentSentinel.INVALID_COMPONENT_MODEL
210 return ComponentSentinel.UNRECOGNIZED_COMPONENT
213class MaskedImageFormatter(ImageFormatter):
214 """A specialized butler interface to FITS archive serialization of
215 the `..MaskedImage` class.
216 """
218 def read_component(
219 self,
220 component: str,
221 tree: Any,
222 archive: FitsInputArchive,
223 ) -> Any:
224 match super().read_component(component, tree, archive):
225 case ComponentSentinel():
226 pass
227 case handled:
228 return handled
229 if not isinstance(tree, MaskedImageSerializationModel):
230 return ComponentSentinel.INVALID_COMPONENT_MODEL
231 match component:
232 case "image":
233 return Image.deserialize(tree.image, archive, bbox=self.pop_bbox_from_parameters())
234 case "mask":
235 return Mask.deserialize(tree.mask, archive, bbox=self.pop_bbox_from_parameters())
236 case "variance":
237 return Image.deserialize(tree.variance, archive, bbox=self.pop_bbox_from_parameters())
238 return ComponentSentinel.UNRECOGNIZED_COMPONENT
241class VisitImageFormatter(MaskedImageFormatter):
242 """A specialized butler interface to FITS archive serialization of
243 the `..VisitImage` class.
244 """
246 def read_component(
247 self,
248 component: str,
249 tree: Any,
250 archive: FitsInputArchive,
251 ) -> Any:
252 match super().read_component(component, tree, archive):
253 case ComponentSentinel():
254 pass
255 case handled:
256 return handled
257 if not isinstance(tree, VisitImageSerializationModel):
258 return ComponentSentinel.INVALID_COMPONENT_MODEL
259 match component:
260 case "psf":
261 return tree.deserialize_psf(archive)
262 return ComponentSentinel.UNRECOGNIZED_COMPONENT