Coverage for tests/test_serialization_streams.py: 96%

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"""Tests for reading archives from in-memory bytes and binary streams.""" 

12 

13from __future__ import annotations 

14 

15import gzip 

16import io 

17import os 

18import tempfile 

19import unittest 

20 

21import numpy as np 

22 

23from lsst.images import Box, Image, Mask 

24from lsst.images.fits import FitsInputArchive 

25from lsst.images.json import JsonInputArchive 

26from lsst.images.serialization import open as open_archive 

27from lsst.images.serialization import read, write 

28 

29try: 

30 import h5py # noqa: F401 -- detect availability for NDF stream skips 

31 

32 from lsst.images.ndf import NdfInputArchive 

33 

34 H5PY_AVAILABLE = True 

35except ImportError: 

36 H5PY_AVAILABLE = False 

37 

38try: 

39 from compression import zstd as _stdlib_zstd # noqa: F401 -- detect zstd availability 

40 

41 ZSTD_AVAILABLE = True 

42except ImportError: 

43 try: 

44 import zstandard # noqa: F401 

45 

46 ZSTD_AVAILABLE = True 

47 except ImportError: 

48 ZSTD_AVAILABLE = False 

49 

50 

51def _zstd_compress(data: bytes) -> bytes: 

52 """Compress with whichever zstd library is available.""" 

53 try: 

54 from compression import zstd 

55 except ImportError: 

56 import zstandard 

57 

58 return zstandard.ZstdCompressor().compress(data) 

59 return zstd.compress(data) 

60 

61 

62def _test_image() -> Image: 

63 return Image(np.arange(16, dtype=np.float32).reshape(4, 4), bbox=Box.factory[0:4, 0:4]) 

64 

65 

66def _serialized_bytes(obj: object, extension: str) -> bytes: 

67 """Write ``obj`` to a temporary file and return the file's content.""" 

68 with tempfile.TemporaryDirectory() as tmp: 

69 path = os.path.join(tmp, f"x{extension}") 

70 write(obj, path) 

71 with open(path, "rb") as f: 

72 return f.read() 

73 

74 

75class FitsOpenTreeStreamTestCase(unittest.TestCase): 

76 """FitsInputArchive.open_tree accepts a seekable binary stream.""" 

77 

78 def test_open_tree_stream(self) -> None: 

79 image = _test_image() 

80 stream = io.BytesIO(_serialized_bytes(image, ".fits")) 

81 with FitsInputArchive.open_tree(stream) as (archive, tree, info): 

82 self.assertEqual(info.schema_name, "image") 

83 result = tree.deserialize(archive) 

84 self.assertIsInstance(result, Image) 

85 np.testing.assert_array_equal(result.array, image.array) 

86 

87 

88class JsonOpenTreeStreamTestCase(unittest.TestCase): 

89 """JsonInputArchive.open_tree accepts a seekable binary stream.""" 

90 

91 def test_open_tree_stream(self) -> None: 

92 image = _test_image() 

93 stream = io.BytesIO(_serialized_bytes(image, ".json")) 

94 with JsonInputArchive.open_tree(stream) as (archive, tree, info): 

95 result = tree.deserialize(archive) 

96 self.assertIsInstance(result, Image) 

97 np.testing.assert_array_equal(result.array, image.array) 

98 

99 

100@unittest.skipUnless(H5PY_AVAILABLE, "h5py not available.") 

101class NdfOpenTreeStreamTestCase(unittest.TestCase): 

102 """NdfInputArchive.open_tree accepts a seekable binary stream.""" 

103 

104 def test_open_tree_stream(self) -> None: 

105 image = _test_image() 

106 stream = io.BytesIO(_serialized_bytes(image, ".sdf")) 

107 with NdfInputArchive.open_tree(stream) as (archive, tree, info): 

108 result = tree.deserialize(archive) 

109 self.assertIsInstance(result, Image) 

110 np.testing.assert_array_equal(result.array, image.array) 

111 

112 

113class ReadStreamTestCase(unittest.TestCase): 

114 """read(io.BytesIO(data)) turns in-memory data into objects.""" 

115 

116 def setUp(self) -> None: 

117 self.image = _test_image() 

118 

119 def test_fits(self) -> None: 

120 result = read(io.BytesIO(_serialized_bytes(self.image, ".fits"))) 

121 self.assertIsInstance(result, Image) 

122 np.testing.assert_array_equal(result.array, self.image.array) 

123 

124 def test_json(self) -> None: 

125 result = read(io.BytesIO(_serialized_bytes(self.image, ".json"))) 

126 self.assertIsInstance(result, Image) 

127 np.testing.assert_array_equal(result.array, self.image.array) 

128 

129 @unittest.skipUnless(H5PY_AVAILABLE, "h5py not available.") 

130 def test_ndf(self) -> None: 

131 result = read(io.BytesIO(_serialized_bytes(self.image, ".sdf"))) 

132 self.assertIsInstance(result, Image) 

133 np.testing.assert_array_equal(result.array, self.image.array) 

134 

135 def test_cls_match_and_mismatch(self) -> None: 

136 data = _serialized_bytes(self.image, ".fits") 

137 result = read(io.BytesIO(data), cls=Image) 

138 self.assertIsInstance(result, Image) 

139 with self.assertRaises(TypeError): 

140 read(io.BytesIO(data), cls=Mask) 

141 

142 def test_kwargs_forwarded(self) -> None: 

143 big = Image(np.arange(64, dtype=np.float32).reshape(8, 8), bbox=Box.factory[0:8, 0:8]) 

144 sub = read(io.BytesIO(_serialized_bytes(big, ".fits")), bbox=Box.factory[2:6, 2:6]) 

145 self.assertEqual(sub.array.shape, (4, 4)) 

146 np.testing.assert_array_equal(sub.array, big.array[2:6, 2:6]) 

147 

148 def test_format_override(self) -> None: 

149 result = read(io.BytesIO(_serialized_bytes(self.image, ".json")), format="json") 

150 self.assertIsInstance(result, Image) 

151 

152 def test_wrong_format_override(self) -> None: 

153 # FITS bytes forced through the JSON backend fail in that backend. 

154 with self.assertRaises(ValueError): 

155 read(io.BytesIO(_serialized_bytes(self.image, ".fits")), format="json") 

156 

157 def test_unrecognized_bytes(self) -> None: 

158 with self.assertRaises(ValueError) as cm: 

159 read(io.BytesIO(b"certainly not a supported format")) 

160 self.assertIn("FITS", str(cm.exception)) 

161 

162 def test_bad_format_name(self) -> None: 

163 with self.assertRaises(ValueError): 

164 read(io.BytesIO(_serialized_bytes(self.image, ".json")), format="asdf") 

165 

166 def test_compressed_stream_raises(self) -> None: 

167 # Compressed streams are the caller's responsibility to decompress; 

168 # the sniff error says what to do. 

169 data = gzip.compress(_serialized_bytes(self.image, ".fits")) 

170 with self.assertRaises(ValueError) as cm: 

171 read(io.BytesIO(data)) 

172 self.assertIn("gzip", str(cm.exception)) 

173 

174 

175class OpenStreamTestCase(unittest.TestCase): 

176 """open() accepts a seekable binary stream, with component reads.""" 

177 

178 DATA_DIR = os.path.join(os.path.dirname(__file__), "data", "schema_v1") 

179 

180 def test_open_stream_components(self) -> None: 

181 visit_image = read(os.path.join(self.DATA_DIR, "visit_image.json")) 

182 stream = io.BytesIO(_serialized_bytes(visit_image, ".fits")) 

183 with open_archive(stream) as reader: 

184 self.assertEqual(reader.info.schema_name, "visit_image") 

185 self.assertIsInstance(reader.metadata, dict) 

186 self.assertIsNotNone(reader.get_component("sky_projection")) 

187 full = reader.read() 

188 self.assertEqual(type(full).__name__, "VisitImage") 

189 

190 

191class CompressedPathTestCase(unittest.TestCase): 

192 """Compressed files read transparently through path-based read().""" 

193 

194 def setUp(self) -> None: 

195 tmp = tempfile.TemporaryDirectory() 

196 self.addCleanup(tmp.cleanup) 

197 self.tmp = tmp.name 

198 self.image = _test_image() 

199 

200 def _write_compressed(self, extension: str, compress) -> str: 

201 data = compress(_serialized_bytes(self.image, extension)) 

202 suffix = ".gz" if compress is gzip.compress else ".zst" 

203 path = os.path.join(self.tmp, f"x{extension}{suffix}") 

204 with open(path, "wb") as f: 

205 f.write(data) 

206 return path 

207 

208 def test_fits_gz(self) -> None: 

209 # Regression: .fits.gz was dispatched to the FITS backend but 

210 # handed to it still compressed. 

211 path = self._write_compressed(".fits", gzip.compress) 

212 result = read(path) 

213 self.assertIsInstance(result, Image) 

214 np.testing.assert_array_equal(result.array, self.image.array) 

215 

216 def test_json_gz(self) -> None: 

217 path = self._write_compressed(".json", gzip.compress) 

218 self.assertIsInstance(read(path), Image) 

219 

220 @unittest.skipUnless(ZSTD_AVAILABLE, "no zstd decompressor available.") 

221 def test_fits_zst(self) -> None: 

222 path = self._write_compressed(".fits", _zstd_compress) 

223 self.assertIsInstance(read(path), Image) 

224 

225 def test_open_fits_gz(self) -> None: 

226 path = self._write_compressed(".fits", gzip.compress) 

227 with open_archive(path) as reader: 

228 self.assertIsInstance(reader.read(), Image) 

229 

230 

231if __name__ == "__main__": 

232 unittest.main()