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