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