Coverage for python / lsst / images / tests / _roundtrip.py: 21%
120 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-18 09:00 +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", "TemporaryButler")
16import tempfile
17import unittest
18import uuid
19from contextlib import ExitStack
20from typing import Any, TypeVar
22import astropy.io.fits
24try:
25 from lsst.daf.butler import Butler, DataCoordinate, DatasetProvenance, DatasetRef, DatasetType
27 HAVE_BUTLER = True
28except ImportError:
29 HAVE_BUTLER = False
31from .. import fits
32from .._generalized_image import GeneralizedImage
33from ..serialization import MetadataValue
35# We need an old-style TypeVar for Sphinx.
36T = TypeVar("T")
39class TemporaryButler:
40 """Make a temporary butler repository.
42 Parameters
43 ----------
44 run
45 Name of a `~lsst.daf.butler.CollectionType.RUN` collection to
46 register and use as the default run for the returned butler.
47 **kwargs
48 A mapping from a dataset type name to its storage class. For each
49 entry, a dataset type will be registered with empty dimensions, and a
50 `~lsst.daf.butler.DatasetRef` will be created and added as an
51 attribute of this class.
53 Raises
54 ------
55 unittest.SkipTest
56 Raised when the context manager is entered if `lsst.daf.butler` could
57 not be imported. This is typically handled by using this context
58 manager within a `unittest.TestCase.subTest` context, which will skip
59 just the butler-required tests in that context while allowing the rest
60 of the test to continue.
61 """
63 def __init__(self, run: str = "test_run", **kwargs: str):
64 self.run = run
65 self._kwargs = kwargs
66 self._exit_stack = ExitStack()
68 def __enter__(self) -> TemporaryButler:
69 if not HAVE_BUTLER:
70 raise unittest.SkipTest("lsst.daf.butler could not be imported.")
71 self._exit_stack.__enter__()
72 root = self._exit_stack.enter_context(
73 tempfile.TemporaryDirectory(ignore_cleanup_errors=True, delete=True)
74 )
75 butler_config = Butler.makeRepo(root)
76 self.butler = self._exit_stack.enter_context(Butler.from_config(butler_config, run=self.run))
77 empty_data_id = DataCoordinate.make_empty(self.butler.dimensions)
78 for name, storage_class in self._kwargs.items():
79 dataset_type = DatasetType(name, self.butler.dimensions.empty, storage_class)
80 try:
81 self.butler.registry.registerDatasetType(dataset_type)
82 except KeyError as err:
83 err.add_note(
84 "Storage class not configured in butler defaults. "
85 "A newer version of daf_butler may be needed."
86 )
87 raise
88 setattr(self, name, DatasetRef(dataset_type, empty_data_id, self.run))
89 return self
91 def __exit__(self, *args: Any) -> bool | None:
92 return self._exit_stack.__exit__(*args)
94 # Just for typing, since this class uses dynamic attributes.
95 def __getattr__(self, name: str) -> DatasetRef:
96 raise AttributeError(name)
99class RoundtripFits[T]:
100 """A context manager for testing FITS-based serialization.
102 Parameters
103 ----------
104 tc
105 A test case object to used for internal checks.
106 original
107 The object to serialize.
108 storage_class
109 A butler storage class name to use. If not provided (or
110 `lsst.daf.butler` cannot be imported), the roundtrip will just use
111 a direct write to a temporary file.
113 Notes
114 -----
115 When entered, this context manager writes the object and reads it back in
116 to the ``result`` attribute. When exited, any temporary files or
117 directories are deleted, but the ``result`` attribute is still usable.
118 In between the `inspect` and `get` methods can be used to perform other
119 tests.
121 This helper internally tests that butler provenance and metadata are saved
122 with any `.GeneralizedImage` object.
123 """
125 def __init__(self, tc: unittest.TestCase, original: T, storage_class: str | None = None):
126 self._original = original
127 self._storage_class = storage_class
128 self._serialized: Any = None
129 self._exit_stack = ExitStack()
130 self._filename: str | None = None
131 self._tc = tc
132 self.result: Any
133 self.butler: Butler | None = None
134 self.ref: DatasetRef | None = None
135 self._test_metadata: dict[str, MetadataValue] = {
136 "roundtrip_test_1": 1,
137 "roundtrip_test_2": 2.5,
138 "roundtrip_test_3": "three",
139 "roundtrip_test_4": True,
140 "roundtrip_test_5": None,
141 }
143 def __enter__(self) -> RoundtripFits[T]:
144 self._exit_stack.__enter__()
145 if isinstance(self._original, GeneralizedImage):
146 self._original.metadata.update(self._test_metadata)
147 if HAVE_BUTLER and self._storage_class is not None:
148 self._run_with_butler()
149 else:
150 self._run_without_butler()
151 if isinstance(self._original, GeneralizedImage):
152 assert isinstance(self.result, GeneralizedImage)
153 for k in self._test_metadata:
154 self._tc.assertEqual(self.result.metadata[k], self._test_metadata[k])
155 del self._original.metadata[k]
156 del self.result.metadata[k]
157 return self
159 def __exit__(self, *args: Any) -> bool | None:
160 return self._exit_stack.__exit__(*args)
162 @property
163 def filename(self) -> str:
164 """The name of the file the object was written to."""
165 if self._filename is None:
166 assert self.butler is not None and self.ref is not None
167 self._filename = self.butler.getURI(self.ref).ospath
168 return self._filename
170 @property
171 def serialized(self) -> Any:
172 """The serialization model for this object
173 (`.serialization.ArchiveTree`).
174 """
175 if self._serialized is None:
176 # The butler code path doesn't give us a way to inspect the
177 # serialized model, so we have to save it again directly to another
178 # file (which we then discard).
179 with tempfile.NamedTemporaryFile(suffix=".fits", delete_on_close=False, delete=True) as tmp:
180 tmp.close()
181 self._serialized = fits.write(self._original, tmp.name)
182 return self._serialized
184 def inspect(self) -> astropy.io.fits.HDUList:
185 """Open the FITS file with Astropy."""
186 return self._exit_stack.enter_context(
187 astropy.io.fits.open(self.filename, disable_image_compression=True)
188 )
190 def get(self, component: str | None = None, storageClass: str | None = None, **kwargs: Any) -> Any:
191 """Perform a partial read.
193 Parameters
194 ----------
195 component
196 Component to read instead of the main object. This requires the
197 roundtrip to use a butler, raising `unittest.SkipTest` otherwise;
198 this generally means these tests should be nested within a
199 `~unittest.TestCase.subTest` context.
200 storageClass
201 Override storage class name to affect the type returned by
202 the get. Only used if a butler is active.
203 **kwargs
204 Keyword arguments either passed directly to `.fits.read` or used
205 as ``parameters`` for a `~lsst.daf.butler.Butler.get`.
207 Return
208 ------
209 object
210 Result of the partial read.
211 """
212 if self.butler is None:
213 if component is not None:
214 raise unittest.SkipTest("Cannot test component reads without a butler.")
215 if storageClass is not None:
216 raise unittest.SkipTest("Cannot test storage class override without a butler")
217 result = fits.read(type(self._original), self.filename, **kwargs)
218 else:
219 assert self.ref is not None, "butler and ref should be None or not together"
220 ref = self.ref
221 if component is not None:
222 ref = ref.makeComponentRef(component)
223 result = self.butler.get(ref, parameters=kwargs, storageClass=storageClass)
224 if isinstance(result, GeneralizedImage):
225 # The metadata the RoundtripFits object added for the test may or
226 # may not be present; strip it if it does so comparisons to the
227 # original are not messed up.
228 for k in self._test_metadata:
229 result.metadata.pop(k, None)
230 return result
232 def _run_with_butler(self) -> None:
233 assert self._storage_class is not None, "Should not use butler if no storage class"
234 butler_helper = self._exit_stack.enter_context(TemporaryButler(test_dataset=self._storage_class))
235 self.butler = butler_helper.butler
236 quantum_id = uuid.uuid4()
237 self.ref = self.butler.put(
238 self._original, butler_helper.test_dataset, provenance=DatasetProvenance(quantum_id=quantum_id)
239 )
240 self.result = self.butler.get(self.ref)
241 if isinstance(self._original, GeneralizedImage):
242 self._tc.assertEqual(
243 DatasetRef.from_simple(self.result.butler_dataset, universe=self.butler.dimensions), self.ref
244 )
245 self._tc.assertEqual(self.result.butler_provenance.quantum_id, quantum_id)
247 def _run_without_butler(self) -> None:
248 tmp = self._exit_stack.enter_context(
249 tempfile.NamedTemporaryFile(suffix=".fits", delete_on_close=False, delete=True)
250 )
251 tmp.close()
252 self._filename = tmp.name
253 self._serialized = fits.write(self._original, tmp.name)
254 read_result = fits.read(type(self._original), tmp.name)
255 self._tc.assertIsNone(read_result.butler_info)
256 self.result = read_result.deserialized