Coverage for python/lsst/images/tests/_roundtrip.py: 85%
156 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-16 15:24 -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__ = ("RoundtripFits", "RoundtripJson", "RoundtripNdf", "TemporaryButler")
16import tempfile
17import uuid
18from abc import ABC, abstractmethod
19from contextlib import ExitStack
20from typing import TYPE_CHECKING, Any, Self, TypeVar
22import astropy.io.fits
23import pytest
24from pydantic_core import from_json
26if TYPE_CHECKING:
27 import h5py
29try:
30 from lsst.daf.butler import Butler, Config, DataCoordinate, DatasetProvenance, DatasetRef, DatasetType
32 HAVE_BUTLER = True
33except ImportError:
34 HAVE_BUTLER = False
36from .._generalized_image import GeneralizedImage
37from ..serialization import MetadataValue, open_archive, read_archive, write_archive
39# We need an old-style TypeVar for Sphinx.
40T = TypeVar("T")
43class TemporaryButler:
44 """Make a temporary butler repository.
46 Parameters
47 ----------
48 run
49 Name of a `~lsst.daf.butler.CollectionType.RUN` collection to
50 register and use as the default run for the returned butler.
51 format
52 Optional on-disk format name (``fits``, ``json``, ``sdf``,
53 ``zarr``, ...) to bind to every storage class registered by
54 ``**kwargs``. When set, the datastore config is overlaid so that
55 `~lsst.images.formatters.GenericFormatter` writes that format for
56 those storage classes, overriding its ``.fits`` default. Leave as
57 `None` to keep the default formatter behaviour.
58 recipe
59 Optional write recipe to bind to every storage class registered by
60 ``**kwargs``.
61 **kwargs
62 A mapping from a dataset type name to its storage class. For each
63 entry, a dataset type will be registered with empty dimensions, and a
64 `~lsst.daf.butler.DatasetRef` will be created and added as an
65 attribute of this class.
67 Notes
68 -----
69 `pytest.skip` is called when the context manager is entered if
70 `lsst.daf.butler` could not be imported, skipping the current test.
71 """
73 def __init__(
74 self,
75 run: str = "test_run",
76 *,
77 format: str | None = None,
78 recipe: str | None = None,
79 **kwargs: str,
80 ) -> None:
81 self.run = run
82 self._format = format
83 self._recipe = recipe
84 self._kwargs = kwargs
85 self._exit_stack = ExitStack()
87 def __enter__(self) -> TemporaryButler:
88 if not HAVE_BUTLER: 88 ↛ 89line 88 didn't jump to line 89 because the condition on line 88 was never true
89 pytest.skip("lsst.daf.butler could not be imported.")
90 self._exit_stack.__enter__()
91 root = self._exit_stack.enter_context(
92 tempfile.TemporaryDirectory(ignore_cleanup_errors=True, delete=True)
93 )
94 write_parameters: dict[str, str] = {}
95 if self._format is not None:
96 write_parameters["format"] = self._format
97 if self._recipe is not None: 97 ↛ 98line 97 didn't jump to line 98 because the condition on line 97 was never true
98 write_parameters["recipe"] = self._recipe
99 if write_parameters:
100 # Overlay a per-storage-class formatter binding so the default
101 # FITS-writing GenericFormatter writes the requested format
102 # instead. Keyed by the storage class name (matched by the
103 # daf_butler formatter factory).
104 overlay = Config(
105 {
106 "datastore": {
107 "formatters": {
108 storage_class: {
109 "formatter": "lsst.images.formatters.GenericFormatter",
110 "parameters": write_parameters,
111 }
112 for storage_class in self._kwargs.values()
113 }
114 }
115 }
116 )
117 butler_config = Butler.makeRepo(root, config=overlay)
118 else:
119 butler_config = Butler.makeRepo(root)
120 self.butler = self._exit_stack.enter_context(Butler.from_config(butler_config, run=self.run))
121 empty_data_id = DataCoordinate.make_empty(self.butler.dimensions)
122 for name, storage_class in self._kwargs.items():
123 dataset_type = DatasetType(name, self.butler.dimensions.empty, storage_class)
124 try:
125 self.butler.registry.registerDatasetType(dataset_type)
126 except KeyError as err:
127 err.add_note(
128 "Storage class not configured in butler defaults. "
129 "A newer version of daf_butler may be needed."
130 )
131 raise
132 setattr(self, name, DatasetRef(dataset_type, empty_data_id, self.run))
133 return self
135 def __exit__(self, *args: Any) -> bool | None:
136 return self._exit_stack.__exit__(*args)
138 # Just for typing, since this class uses dynamic attributes.
139 def __getattr__(self, name: str) -> DatasetRef:
140 raise AttributeError(name)
143class RoundtripBase[T](ABC):
144 """A context manager for testing serialization.
146 Parameters
147 ----------
148 original
149 The object to serialize.
150 storage_class
151 A butler storage class name to use. If not provided (or
152 `lsst.daf.butler` cannot be imported), the roundtrip will just use
153 a direct write to a temporary file.
154 recipe
155 Write recipe used to control butler puts; only used when roundtripping
156 through a butler.
157 **kwargs
158 Keyword arguments to pass to `write`, usually equivalent to what
159 ``recipe`` resolves to; ignored when roundtripping through a butler.
161 Notes
162 -----
163 When entered, this context manager writes the object and reads it back in
164 to the ``result`` attribute. When exited, any temporary files or
165 directories are deleted, but the ``result`` attribute is still usable.
166 In between the `inspect` and `get` methods can be used to perform other
167 tests.
169 This helper internally tests that butler provenance and metadata are saved
170 with any `.GeneralizedImage` object.
171 """
173 def __init__(
174 self,
175 original: T,
176 storage_class: str | None = None,
177 recipe: str | None = None,
178 **kwargs: Any,
179 ) -> None:
180 self._original = original
181 self._storage_class = storage_class
182 self._serialized: Any = None
183 self._exit_stack = ExitStack()
184 self._filename: str | None = None
185 self._recipe = recipe
186 self._write_kwargs = kwargs
187 self.result: Any
188 self.butler: Butler | None = None
189 self.ref: DatasetRef | None = None
190 self._test_metadata: dict[str, MetadataValue] = {
191 "roundtrip_test_1": 1,
192 "roundtrip_test_2": 2.5,
193 "roundtrip_test_3": "three",
194 "roundtrip_test_4": True,
195 "roundtrip_test_5": None,
196 }
198 def __enter__(self) -> Self:
199 self._exit_stack.__enter__()
200 if isinstance(self._original, GeneralizedImage):
201 self._original.metadata.update(self._test_metadata)
202 if HAVE_BUTLER and self._storage_class is not None:
203 self._run_with_butler()
204 else:
205 self._run_without_butler()
206 if isinstance(self._original, GeneralizedImage):
207 assert isinstance(self.result, GeneralizedImage)
208 for k in self._test_metadata:
209 assert self.result.metadata[k] == self._test_metadata[k]
210 del self._original.metadata[k]
211 del self.result.metadata[k]
212 return self
214 def __exit__(self, *args: Any) -> bool | None:
215 return self._exit_stack.__exit__(*args)
217 @property
218 def filename(self) -> str:
219 """The name of the file the object was written to."""
220 if self._filename is None:
221 assert self.butler is not None and self.ref is not None
222 self._filename = self.butler.getURI(self.ref).ospath
223 return self._filename
225 @property
226 def serialized(self) -> Any:
227 """The serialization model for this object
228 (`.serialization.ArchiveTree`).
229 """
230 if self._serialized is None:
231 # The butler code path doesn't give us a way to inspect the
232 # serialized model, so we have to save it again directly to another
233 # file (which we then discard).
234 with tempfile.NamedTemporaryFile(
235 suffix=self._get_extension(), delete_on_close=False, delete=True
236 ) as tmp:
237 tmp.close()
238 self._serialized = write_archive(self._original, tmp.name)
239 return self._serialized
241 def get(self, component: str | None = None, storageClass: str | None = None, **kwargs: Any) -> Any:
242 """Perform a partial read.
244 Parameters
245 ----------
246 component
247 Component to read instead of the main object. This requires the
248 roundtrip to use a butler; `pytest.skip` is called otherwise.
249 Place calls to this method in a dedicated test function that
250 contains only component-read assertions, so the skip does not
251 suppress unrelated always-run assertions in other test functions.
252 storageClass
253 Override storage class name to affect the type returned by
254 the get. Only used if a butler is active.
255 **kwargs
256 Keyword arguments either passed directly to
257 `~lsst.images.serialization.read_archive` or used as
258 ``parameters`` for a `~lsst.daf.butler.Butler.get`.
260 Return
261 ------
262 object
263 Result of the partial read.
264 """
265 if self.butler is None: 265 ↛ 266line 265 didn't jump to line 266 because the condition on line 265 was never true
266 if component is not None:
267 pytest.skip("Cannot test component reads without a butler.")
268 if storageClass is not None:
269 pytest.skip("Cannot test storage class override without a butler")
270 result = read_archive(self.filename, type(self._original), **kwargs)
271 else:
272 assert self.ref is not None, "butler and ref should be None or not together"
273 ref = self.ref
274 if component is not None:
275 ref = ref.makeComponentRef(component)
276 result = self.butler.get(ref, parameters=kwargs, storageClass=storageClass)
277 if isinstance(result, GeneralizedImage):
278 # The metadata the RoundtripFits object added for the test may or
279 # may not be present; strip it if it does so comparisons to the
280 # original are not messed up.
281 for k in self._test_metadata:
282 result.metadata.pop(k, None)
283 if component == "components" and isinstance(result, dict):
284 # A special case component that returns a dict of components
285 # that each need to have their metadata potentially cleaned up.
286 for value in result.values():
287 if isinstance(value, GeneralizedImage):
288 for k in self._test_metadata:
289 value.metadata.pop(k, None)
290 return result
292 def _run_with_butler(self) -> None:
293 assert self._storage_class is not None, "Should not use butler if no storage class"
294 # ``GenericFormatter`` defaults to FITS; tell the temporary butler
295 # which format this Roundtrip variant wants so the on-disk file
296 # matches ``_get_extension()`` on the round-trip check below.
297 fmt = self._get_extension().lstrip(".")
298 butler_helper = self._exit_stack.enter_context(
299 TemporaryButler(test_dataset=self._storage_class, format=fmt, recipe=self._recipe)
300 )
301 self.butler = butler_helper.butler
302 quantum_id = uuid.uuid4()
303 self.ref = self.butler.put(
304 self._original, butler_helper.test_dataset, provenance=DatasetProvenance(quantum_id=quantum_id)
305 )
306 self.result = self.butler.get(self.ref)
307 if isinstance(self._original, GeneralizedImage): 307 ↛ 313line 307 didn't jump to line 313 because the condition on line 307 was always true
308 assert (
309 DatasetRef.from_simple(self.result.butler_dataset, universe=self.butler.dimensions)
310 == self.ref
311 )
312 assert self.result.butler_provenance.quantum_id == quantum_id
313 assert self.filename.endswith(self._get_extension()), (
314 f"{self.filename} did not end with {self._get_extension()}"
315 )
317 def _run_without_butler(self) -> None:
318 tmp = self._exit_stack.enter_context(
319 tempfile.NamedTemporaryFile(suffix=self._get_extension(), delete_on_close=False, delete=True)
320 )
321 tmp.close()
322 self._filename = tmp.name
323 self._serialized = write_archive(self._original, tmp.name, **self._write_kwargs)
324 with open_archive(tmp.name, type(self._original)) as reader:
325 assert reader.butler_info is None
326 self.result = reader.read()
328 @abstractmethod
329 def _get_extension(self) -> str:
330 raise NotImplementedError()
333class RoundtripFits[T](RoundtripBase[T]):
334 def inspect(self) -> astropy.io.fits.HDUList:
335 """Open the FITS file with Astropy."""
336 return self._exit_stack.enter_context(
337 astropy.io.fits.open(self.filename, disable_image_compression=True)
338 )
340 def _get_extension(self) -> str:
341 return ".fits"
344class RoundtripJson[T](RoundtripBase[T]):
345 def inspect(self) -> dict[str, Any]:
346 """Read the JSON file as a dictionary."""
347 with open(self.filename, "rb") as stream:
348 return from_json(stream.read())
350 def _get_extension(self) -> str:
351 return ".json"
354class RoundtripNdf[T](RoundtripBase[T]):
355 def inspect(self) -> h5py.File:
356 """Open the NDF file with h5py."""
357 import h5py
359 return self._exit_stack.enter_context(h5py.File(self.filename, "r"))
361 def _get_extension(self) -> str:
362 return ".sdf"