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