Coverage for python/lsst/images/tests/_roundtrip.py: 86%
169 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-29 22:40 +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 .. import fits, json
37from .._generalized_image import GeneralizedImage
38from ..serialization import ArchiveTree, MetadataValue
39from ..serialization import open as open_archive
40from ..serialization import read as read_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(suffix=".fits", delete_on_close=False, delete=True) as tmp:
246 tmp.close()
247 self._serialized = fits.write(self._original, tmp.name)
248 return self._serialized
250 def get(self, component: str | None = None, storageClass: str | None = None, **kwargs: Any) -> Any:
251 """Perform a partial read.
253 Parameters
254 ----------
255 component
256 Component to read instead of the main object. This requires the
257 roundtrip to use a butler, raising `unittest.SkipTest` otherwise;
258 this generally means these tests should be nested within a
259 `~unittest.TestCase.subTest` context.
260 storageClass
261 Override storage class name to affect the type returned by
262 the get. Only used if a butler is active.
263 **kwargs
264 Keyword arguments either passed directly to
265 `~lsst.images.serialization.read` or used as ``parameters`` for a
266 `~lsst.daf.butler.Butler.get`.
268 Return
269 ------
270 object
271 Result of the partial read.
272 """
273 if self.butler is None: 273 ↛ 274line 273 didn't jump to line 274 because the condition on line 273 was never true
274 if component is not None:
275 raise unittest.SkipTest("Cannot test component reads without a butler.")
276 if storageClass is not None:
277 raise unittest.SkipTest("Cannot test storage class override without a butler")
278 result = read_archive(self.filename, type(self._original), **kwargs)
279 else:
280 assert self.ref is not None, "butler and ref should be None or not together"
281 ref = self.ref
282 if component is not None:
283 ref = ref.makeComponentRef(component)
284 result = self.butler.get(ref, parameters=kwargs, storageClass=storageClass)
285 if isinstance(result, GeneralizedImage):
286 # The metadata the RoundtripFits object added for the test may or
287 # may not be present; strip it if it does so comparisons to the
288 # original are not messed up.
289 for k in self._test_metadata:
290 result.metadata.pop(k, None)
291 if component == "components" and isinstance(result, dict):
292 # A special case component that returns a dict of components
293 # that each need to have their metadata potentially cleaned up.
294 for value in result.values():
295 if isinstance(value, GeneralizedImage):
296 for k in self._test_metadata:
297 value.metadata.pop(k, None)
298 return result
300 def _run_with_butler(self) -> None:
301 assert self._storage_class is not None, "Should not use butler if no storage class"
302 # ``GenericFormatter`` defaults to FITS; tell the temporary butler
303 # which format this Roundtrip variant wants so the on-disk file
304 # matches ``_get_extension()`` on the round-trip check below.
305 fmt = self._get_extension().lstrip(".")
306 butler_helper = self._exit_stack.enter_context(
307 TemporaryButler(test_dataset=self._storage_class, format=fmt, recipe=self._recipe)
308 )
309 self.butler = butler_helper.butler
310 quantum_id = uuid.uuid4()
311 self.ref = self.butler.put(
312 self._original, butler_helper.test_dataset, provenance=DatasetProvenance(quantum_id=quantum_id)
313 )
314 self.result = self.butler.get(self.ref)
315 if isinstance(self._original, GeneralizedImage): 315 ↛ 320line 315 didn't jump to line 320 because the condition on line 315 was always true
316 self._tc.assertEqual(
317 DatasetRef.from_simple(self.result.butler_dataset, universe=self.butler.dimensions), self.ref
318 )
319 self._tc.assertEqual(self.result.butler_provenance.quantum_id, quantum_id)
320 self._tc.assertTrue(
321 self.filename.endswith(self._get_extension()),
322 f"{self.filename} did not end with {self._get_extension()}",
323 )
325 def _run_without_butler(self) -> None:
326 tmp = self._exit_stack.enter_context(
327 tempfile.NamedTemporaryFile(suffix=self._get_extension(), delete_on_close=False, delete=True)
328 )
329 tmp.close()
330 self._filename = tmp.name
331 self._serialized = self._write(self._original, tmp.name, **self._write_kwargs)
332 with open_archive(tmp.name, type(self._original)) as reader:
333 self._tc.assertIsNone(reader.butler_info)
334 self.result = reader.read()
336 @abstractmethod
337 def _get_extension(self) -> str:
338 raise NotImplementedError()
340 @abstractmethod
341 def _write(self, obj: Any, filename: str) -> ArchiveTree:
342 raise NotImplementedError()
345class RoundtripFits[T](RoundtripBase[T]):
346 def inspect(self) -> astropy.io.fits.HDUList:
347 """Open the FITS file with Astropy."""
348 return self._exit_stack.enter_context(
349 astropy.io.fits.open(self.filename, disable_image_compression=True)
350 )
352 def _get_extension(self) -> str:
353 return ".fits"
355 def _write(self, obj: Any, filename: str) -> ArchiveTree:
356 return fits.write(obj, filename)
359class RoundtripJson[T](RoundtripBase[T]):
360 def inspect(self) -> dict[str, Any]:
361 """Read the JSON file as a dictionary."""
362 with open(self.filename, "rb") as stream:
363 return from_json(stream.read())
365 def _get_extension(self) -> str:
366 return ".json"
368 def _write(self, obj: Any, filename: str) -> ArchiveTree:
369 return json.write(obj, filename)
372class RoundtripNdf[T](RoundtripBase[T]):
373 def inspect(self) -> h5py.File:
374 """Open the NDF file with h5py."""
375 import h5py
377 return self._exit_stack.enter_context(h5py.File(self.filename, "r"))
379 def _get_extension(self) -> str:
380 return ".sdf"
382 def _write(self, obj: Any, filename: str) -> ArchiveTree:
383 from .. import ndf
385 return ndf.write(obj, filename)