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_archive, read_archive, write_archive 

27 

28try: 

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

30 

31 from lsst.images.ndf import NdfInputArchive 

32 

33 H5PY_AVAILABLE = True 

34except ImportError: 

35 H5PY_AVAILABLE = False 

36 

37try: 

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

39 

40 ZSTD_AVAILABLE = True 

41except ImportError: 

42 try: 

43 import zstandard # noqa: F401 

44 

45 ZSTD_AVAILABLE = True 

46 except ImportError: 

47 ZSTD_AVAILABLE = False 

48 

49 

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

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

52 try: 

53 from compression import zstd 

54 except ImportError: 

55 import zstandard 

56 

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

58 return zstd.compress(data) 

59 

60 

61def _test_image() -> Image: 

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

63 

64 

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

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

67 with tempfile.TemporaryDirectory() as tmp: 

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

69 write_archive(obj, path) 

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

71 return f.read() 

72 

73 

74class FitsOpenTreeStreamTestCase(unittest.TestCase): 

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

76 

77 def test_open_tree_stream(self) -> None: 

78 image = _test_image() 

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

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

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

82 result = tree.deserialize(archive) 

83 self.assertIsInstance(result, Image) 

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

85 

86 

87class JsonOpenTreeStreamTestCase(unittest.TestCase): 

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

89 

90 def test_open_tree_stream(self) -> None: 

91 image = _test_image() 

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

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

94 result = tree.deserialize(archive) 

95 self.assertIsInstance(result, Image) 

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

97 

98 

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

100class NdfOpenTreeStreamTestCase(unittest.TestCase): 

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

102 

103 def test_open_tree_stream(self) -> None: 

104 image = _test_image() 

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

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

107 result = tree.deserialize(archive) 

108 self.assertIsInstance(result, Image) 

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

110 

111 

112class ReadStreamTestCase(unittest.TestCase): 

113 """read_archive(io.BytesIO(data)) turns in-memory data into objects.""" 

114 

115 def setUp(self) -> None: 

116 self.image = _test_image() 

117 

118 def test_fits(self) -> None: 

119 result = read_archive(io.BytesIO(_serialized_bytes(self.image, ".fits"))) 

120 self.assertIsInstance(result, Image) 

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

122 

123 def test_json(self) -> None: 

124 result = read_archive(io.BytesIO(_serialized_bytes(self.image, ".json"))) 

125 self.assertIsInstance(result, Image) 

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

127 

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

129 def test_ndf(self) -> None: 

130 result = read_archive(io.BytesIO(_serialized_bytes(self.image, ".sdf"))) 

131 self.assertIsInstance(result, Image) 

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

133 

134 def test_cls_match_and_mismatch(self) -> None: 

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

136 result = read_archive(io.BytesIO(data), cls=Image) 

137 self.assertIsInstance(result, Image) 

138 with self.assertRaises(TypeError): 

139 read_archive(io.BytesIO(data), cls=Mask) 

140 

141 def test_kwargs_forwarded(self) -> None: 

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

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

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

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

146 

147 def test_format_override(self) -> None: 

148 result = read_archive(io.BytesIO(_serialized_bytes(self.image, ".json")), format="json") 

149 self.assertIsInstance(result, Image) 

150 

151 def test_wrong_format_override(self) -> None: 

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

153 with self.assertRaises(ValueError): 

154 read_archive(io.BytesIO(_serialized_bytes(self.image, ".fits")), format="json") 

155 

156 def test_unrecognized_bytes(self) -> None: 

157 with self.assertRaises(ValueError) as cm: 

158 read_archive(io.BytesIO(b"certainly not a supported format")) 

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

160 

161 def test_bad_format_name(self) -> None: 

162 with self.assertRaises(ValueError): 

163 read_archive(io.BytesIO(_serialized_bytes(self.image, ".json")), format="asdf") 

164 

165 def test_compressed_stream_raises(self) -> None: 

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

167 # the sniff error says what to do. 

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

169 with self.assertRaises(ValueError) as cm: 

170 read_archive(io.BytesIO(data)) 

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

172 

173 

174class OpenStreamTestCase(unittest.TestCase): 

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

176 

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

178 

179 def test_open_stream_components(self) -> None: 

180 visit_image = read_archive(os.path.join(self.DATA_DIR, "visit_image.json")) 

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

182 with open_archive(stream) as reader: 

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

184 self.assertIsInstance(reader.metadata, dict) 

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

186 full = reader.read() 

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

188 

189 

190class CompressedPathTestCase(unittest.TestCase): 

191 """Compressed files read transparently via path-based read_archive().""" 

192 

193 def setUp(self) -> None: 

194 tmp = tempfile.TemporaryDirectory() 

195 self.addCleanup(tmp.cleanup) 

196 self.tmp = tmp.name 

197 self.image = _test_image() 

198 

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

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

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

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

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

204 f.write(data) 

205 return path 

206 

207 def test_fits_gz(self) -> None: 

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

209 # handed to it still compressed. 

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

211 result = read_archive(path) 

212 self.assertIsInstance(result, Image) 

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

214 

215 def test_json_gz(self) -> None: 

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

217 self.assertIsInstance(read_archive(path), Image) 

218 

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

220 def test_fits_zst(self) -> None: 

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

222 self.assertIsInstance(read_archive(path), Image) 

223 

224 def test_open_fits_gz(self) -> None: 

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

226 with open_archive(path) as reader: 

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

228 

229 

230if __name__ == "__main__": 

231 unittest.main()