Coverage for python/lsst/images/tests/_roundtrip.py: 85%
157 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-22 10:09 +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__ = ("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
222 assert self.ref is not None
223 self._filename = self.butler.getURI(self.ref).ospath
224 return self._filename
226 @property
227 def serialized(self) -> Any:
228 """The serialization model for this object
229 (`.serialization.ArchiveTree`).
230 """
231 if self._serialized is None:
232 # The butler code path doesn't give us a way to inspect the
233 # serialized model, so we have to save it again directly to another
234 # file (which we then discard).
235 with tempfile.NamedTemporaryFile(
236 suffix=self._get_extension(), delete_on_close=False, delete=True
237 ) as tmp:
238 tmp.close()
239 self._serialized = write_archive(self._original, tmp.name)
240 return self._serialized
242 def get(self, component: str | None = None, storageClass: str | None = None, **kwargs: Any) -> Any:
243 """Perform a partial read.
245 Parameters
246 ----------
247 component
248 Component to read instead of the main object. This requires the
249 roundtrip to use a butler; `pytest.skip` is called otherwise.
250 Place calls to this method in a dedicated test function that
251 contains only component-read assertions, so the skip does not
252 suppress unrelated always-run assertions in other test functions.
253 storageClass
254 Override storage class name to affect the type returned by
255 the get. Only used if a butler is active.
256 **kwargs
257 Keyword arguments either passed directly to
258 `~lsst.images.serialization.read_archive` or used as
259 ``parameters`` for a `~lsst.daf.butler.Butler.get`.
261 Returns
262 -------
263 object
264 Result of the partial read.
265 """
266 if self.butler is None: 266 ↛ 267line 266 didn't jump to line 267 because the condition on line 266 was never true
267 if component is not None:
268 pytest.skip("Cannot test component reads without a butler.")
269 if storageClass is not None:
270 pytest.skip("Cannot test storage class override without a butler")
271 result = read_archive(self.filename, type(self._original), **kwargs)
272 else:
273 assert self.ref is not None, "butler and ref should be None or not together"
274 ref = self.ref
275 if component is not None:
276 ref = ref.makeComponentRef(component)
277 result = self.butler.get(ref, parameters=kwargs, storageClass=storageClass)
278 if isinstance(result, GeneralizedImage):
279 # The metadata the RoundtripFits object added for the test may or
280 # may not be present; strip it if it does so comparisons to the
281 # original are not messed up.
282 for k in self._test_metadata:
283 result.metadata.pop(k, None)
284 if component == "components" and isinstance(result, dict):
285 # A special case component that returns a dict of components
286 # that each need to have their metadata potentially cleaned up.
287 for value in result.values():
288 if isinstance(value, GeneralizedImage):
289 for k in self._test_metadata:
290 value.metadata.pop(k, None)
291 return result
293 def _run_with_butler(self) -> None:
294 assert self._storage_class is not None, "Should not use butler if no storage class"
295 # ``GenericFormatter`` defaults to FITS; tell the temporary butler
296 # which format this Roundtrip variant wants so the on-disk file
297 # matches ``_get_extension()`` on the round-trip check below.
298 fmt = self._get_extension().lstrip(".")
299 butler_helper = self._exit_stack.enter_context(
300 TemporaryButler(test_dataset=self._storage_class, format=fmt, recipe=self._recipe)
301 )
302 self.butler = butler_helper.butler
303 quantum_id = uuid.uuid4()
304 self.ref = self.butler.put(
305 self._original, butler_helper.test_dataset, provenance=DatasetProvenance(quantum_id=quantum_id)
306 )
307 self.result = self.butler.get(self.ref)
308 if isinstance(self._original, GeneralizedImage): 308 ↛ 314line 308 didn't jump to line 314 because the condition on line 308 was always true
309 assert (
310 DatasetRef.from_simple(self.result.butler_dataset, universe=self.butler.dimensions)
311 == self.ref
312 )
313 assert self.result.butler_provenance.quantum_id == quantum_id
314 assert self.filename.endswith(self._get_extension()), (
315 f"{self.filename} did not end with {self._get_extension()}"
316 )
318 def _run_without_butler(self) -> None:
319 tmp = self._exit_stack.enter_context(
320 tempfile.NamedTemporaryFile(suffix=self._get_extension(), delete_on_close=False, delete=True)
321 )
322 tmp.close()
323 self._filename = tmp.name
324 self._serialized = write_archive(self._original, tmp.name, **self._write_kwargs)
325 with open_archive(tmp.name, type(self._original)) as reader:
326 assert reader.butler_info is None
327 self.result = reader.read()
329 @abstractmethod
330 def _get_extension(self) -> str:
331 raise NotImplementedError()
334class RoundtripFits[T](RoundtripBase[T]):
335 def inspect(self) -> astropy.io.fits.HDUList:
336 """Open the FITS file with Astropy."""
337 return self._exit_stack.enter_context(
338 astropy.io.fits.open(self.filename, disable_image_compression=True)
339 )
341 def _get_extension(self) -> str:
342 return ".fits"
345class RoundtripJson[T](RoundtripBase[T]):
346 def inspect(self) -> dict[str, Any]:
347 """Read the JSON file as a dictionary."""
348 with open(self.filename, "rb") as stream:
349 return from_json(stream.read())
351 def _get_extension(self) -> str:
352 return ".json"
355class RoundtripNdf[T](RoundtripBase[T]):
356 def inspect(self) -> h5py.File:
357 """Open the NDF file with h5py."""
358 import h5py
360 return self._exit_stack.enter_context(h5py.File(self.filename, "r"))
362 def _get_extension(self) -> str:
363 return ".sdf"