Coverage for python / lsst / images / tests / _roundtrip.py: 26%

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. 

11 

12from __future__ import annotations 

13 

14__all__ = ("RoundtripFits", "TemporaryButler") 

15 

16import tempfile 

17import unittest 

18from contextlib import ExitStack 

19from typing import Any, TypeVar 

20 

21import astropy.io.fits 

22 

23try: 

24 from lsst.daf.butler import Butler, DataCoordinate, DatasetRef, DatasetType 

25 

26 HAVE_BUTLER = True 

27except ImportError: 

28 HAVE_BUTLER = False 

29 

30from .. import fits 

31 

32# We need an old-style TypeVar for Sphinx. 

33T = TypeVar("T") 

34 

35 

36class TemporaryButler: 

37 """Make a temporary butler repository. 

38 

39 Parameters 

40 ---------- 

41 run 

42 Name of a `~lsst.daf.butler.CollectionType.RUN` collection to 

43 register and use as the default run for the returned butler. 

44 **kwargs 

45 A mapping from a dataset type name to its storage class. For each 

46 entry, a dataset type will be registered with empty dimensions, and a 

47 `~lsst.daf.butler.DatasetRef` will be created and added as an 

48 attribute of this class. 

49 

50 Raises 

51 ------ 

52 unittest.SkipTest 

53 Raised when the context manager is entered if `lsst.daf.butler` could 

54 not be imported. This is typically handled by using this context 

55 manager within a `unittest.TestCase.subTest` context, which will skip 

56 just the butler-required tests in that context while allowing the rest 

57 of the test to continue. 

58 """ 

59 

60 def __init__(self, run: str = "test_run", **kwargs: str): 

61 self.run = run 

62 self._kwargs = kwargs 

63 self._exit_stack = ExitStack() 

64 

65 def __enter__(self) -> TemporaryButler: 

66 if not HAVE_BUTLER: 

67 raise unittest.SkipTest("lsst.daf.butler could not be imported.") 

68 self._exit_stack.__enter__() 

69 root = self._exit_stack.enter_context( 

70 tempfile.TemporaryDirectory(ignore_cleanup_errors=True, delete=True) 

71 ) 

72 butler_config = Butler.makeRepo(root) 

73 self.butler = self._exit_stack.enter_context(Butler.from_config(butler_config, run=self.run)) 

74 empty_data_id = DataCoordinate.make_empty(self.butler.dimensions) 

75 for name, storage_class in self._kwargs.items(): 

76 dataset_type = DatasetType(name, self.butler.dimensions.empty, storage_class) 

77 try: 

78 self.butler.registry.registerDatasetType(dataset_type) 

79 except KeyError as err: 

80 err.add_note( 

81 "Storage class not configured in butler defaults. " 

82 "A newer version of daf_butler may be needed." 

83 ) 

84 raise 

85 setattr(self, name, DatasetRef(dataset_type, empty_data_id, self.run)) 

86 return self 

87 

88 def __exit__(self, *args: Any) -> bool | None: 

89 return self._exit_stack.__exit__(*args) 

90 

91 # Just for typing, since this class uses dynamic attributes. 

92 def __getattr__(self, name: str) -> DatasetRef: 

93 raise AttributeError(name) 

94 

95 

96class RoundtripFits[T]: 

97 """A context manager for testing FITS-based serialization. 

98 

99 Parameters 

100 ---------- 

101 tc 

102 A test case object to used for internal checks. 

103 original 

104 The object to serialize. 

105 storage_class 

106 A butler storage class name to use. If not provided (or 

107 `lsst.daf.butler` cannot be imported), the roundtrip will just use 

108 a direct write to a temporary file. 

109 

110 Notes 

111 ----- 

112 When entered, this context manager writes the object and reads it back in 

113 to the ``result`` attribute. When exited, any temporary files or 

114 directories are deleted, but the ``result`` attribute is still usable. 

115 In between the `inspect` and `get` methods can be used to perform other 

116 tests. 

117 """ 

118 

119 def __init__(self, tc: unittest.TestCase, original: T, storage_class: str | None = None): 

120 self._original = original 

121 self._storage_class = storage_class 

122 self._serialized: Any = None 

123 self._exit_stack = ExitStack() 

124 self._filename: str | None = None 

125 self.result: Any 

126 self.butler: Butler | None = None 

127 self.ref: DatasetRef | None = None 

128 

129 def __enter__(self) -> RoundtripFits[T]: 

130 self._exit_stack.__enter__() 

131 if HAVE_BUTLER and self._storage_class is not None: 

132 self._run_with_butler() 

133 else: 

134 self._run_without_butler() 

135 return self 

136 

137 def __exit__(self, *args: Any) -> bool | None: 

138 return self._exit_stack.__exit__(*args) 

139 

140 @property 

141 def filename(self) -> str: 

142 """The name of the file the object was written to.""" 

143 if self._filename is None: 

144 assert self.butler is not None and self.ref is not None 

145 self._filename = self.butler.getURI(self.ref).ospath 

146 return self._filename 

147 

148 @property 

149 def serialized(self) -> Any: 

150 """The serialization model for this object 

151 (`.serialization.ArchiveTree`). 

152 """ 

153 if self._serialized is None: 

154 # The butler code path doesn't give us a way to inspect the 

155 # serialized model, so we have to save it again directly to another 

156 # file (which we then discard). 

157 with tempfile.NamedTemporaryFile(suffix=".fits", delete_on_close=False, delete=True) as tmp: 

158 tmp.close() 

159 self._serialized = fits.write(self._original, tmp.name) 

160 return self._serialized 

161 

162 def inspect(self) -> astropy.io.fits.HDUList: 

163 """Open the FITS file with Astropy.""" 

164 return self._exit_stack.enter_context( 

165 astropy.io.fits.open(self.filename, disable_image_compression=True) 

166 ) 

167 

168 def get(self, component: str | None = None, **kwargs: Any) -> Any: 

169 """Perform a partial read. 

170 

171 Parameters 

172 ---------- 

173 component 

174 Component to read instead of the main object. This requires the 

175 roundtrip to use a butler, raising `unittest.SkipTest` otherwise; 

176 this generally means these tests should be nested within a 

177 `~unittest.TestCase.subTest` context. 

178 **kwargs 

179 Keyword arguments either passed directly to `.fits.read` or used 

180 as ``parameters`` for a `~lsst.daf.butler.Butler.get`. 

181 

182 Return 

183 ------ 

184 object 

185 Result of the partial read. 

186 """ 

187 if self.butler is None: 

188 if component is not None: 

189 raise unittest.SkipTest("Cannot test component reads without a butler.") 

190 return fits.read(type(self._original), self.filename, **kwargs) 

191 else: 

192 assert self.ref is not None, "butler and ref should be None or not together" 

193 ref = self.ref 

194 if component is not None: 

195 ref = ref.makeComponentRef(component) 

196 return self.butler.get(ref, parameters=kwargs) 

197 

198 def _run_with_butler(self) -> None: 

199 assert self._storage_class is not None, "Should not use butler if no storage class" 

200 butler_helper = self._exit_stack.enter_context(TemporaryButler(test_dataset=self._storage_class)) 

201 self.butler = butler_helper.butler 

202 self.ref = self.butler.put(self._original, butler_helper.test_dataset) 

203 self.result = self.butler.get(self.ref) 

204 

205 def _run_without_butler(self) -> None: 

206 tmp = self._exit_stack.enter_context( 

207 tempfile.NamedTemporaryFile(suffix=".fits", delete_on_close=False, delete=True) 

208 ) 

209 tmp.close() 

210 self._filename = tmp.name 

211 self._serialized = fits.write(self._original, tmp.name) 

212 self.result = fits.read(type(self._original), tmp.name)