Coverage for python/lsst/images/serialization/_backends.py: 99%

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__ = ("Backend", "backend_for_name", "backend_for_path", "backend_for_stream") 

15 

16import contextlib 

17import dataclasses 

18import gzip 

19import os 

20import shutil 

21import tempfile 

22from collections.abc import Callable 

23from typing import IO, TYPE_CHECKING, cast 

24 

25from lsst.resources import ResourcePath, ResourcePathExpression 

26 

27if TYPE_CHECKING: 

28 from typing_extensions import TypeIs 

29 

30 from ._input_archive import InputArchive 

31 

32_GZIP_MAGIC = b"\x1f\x8b" 

33"""Leading bytes of a gzip member (diagnostics only; compressed streams 

34are the caller's responsibility to decompress).""" 

35 

36_ZSTD_MAGIC = b"\x28\xb5\x2f\xfd" 

37"""Leading bytes of a zstd frame (diagnostics only; compressed streams 

38are the caller's responsibility to decompress).""" 

39 

40_COMPRESSION_SUFFIXES = (".gz", ".zst") 

41"""File name suffixes implying whole-file compression.""" 

42 

43_FITS_MAGIC = b"SIMPLE =" 

44"""Leading bytes of a standard FITS primary header.""" 

45 

46_HDF5_MAGIC = b"\x89HDF\r\n\x1a\n" 

47"""Leading bytes of an HDF5 file.""" 

48 

49 

50def _is_binary_stream(obj: object) -> TypeIs[IO[bytes]]: 

51 """Return whether ``obj`` is a readable, seekable binary stream. 

52 

53 ``seek`` is the discriminator against `~lsst.resources.ResourcePath`, 

54 which also has a ``read`` method. 

55 """ 

56 return hasattr(obj, "read") and hasattr(obj, "seek") 

57 

58 

59def _path_is_compressed(path: ResourcePathExpression) -> bool: 

60 """Return whether ``path``'s file name ends in a compression suffix. 

61 

62 Parameters 

63 ---------- 

64 path 

65 Path to inspect; convertible to `lsst.resources.ResourcePath`. 

66 """ 

67 return ResourcePath(path).basename().endswith(_COMPRESSION_SUFFIXES) 

68 

69 

70def _open_zstd_stream(raw: IO[bytes]) -> IO[bytes]: 

71 """Wrap ``raw`` in a streaming zstd decompressor, preferring the stdlib. 

72 

73 The imports are function-scoped optional-dependency guards: 

74 ``compression.zstd`` exists only on Python >= 3.14 and ``zstandard`` 

75 is a third-party fallback that is not a required dependency. 

76 

77 Parameters 

78 ---------- 

79 raw 

80 Binary stream positioned at the start of a zstd frame. 

81 

82 Raises 

83 ------ 

84 ValueError 

85 If no zstd decompressor is available. 

86 """ 

87 try: 

88 from compression import zstd 

89 except ImportError: 

90 try: 

91 import zstandard 

92 except ImportError: 

93 raise ValueError( 

94 "Data is zstd-compressed, but no zstd decompressor is " 

95 "available; install the 'zstandard' package or use " 

96 "Python >= 3.14." 

97 ) from None 

98 return cast(IO[bytes], zstandard.ZstdDecompressor().stream_reader(raw)) 

99 return cast(IO[bytes], zstd.ZstdFile(raw)) 

100 

101 

102def _decompress_path_to_temp_file(path: ResourcePathExpression) -> IO[bytes]: 

103 """Stream-decompress a compressed file into an anonymous temporary file. 

104 

105 The decompressor is selected from the file name's compression suffix. 

106 A compressed file may be arbitrarily large, so its decompressed 

107 content goes to disk in bounded-memory chunks rather than into memory. 

108 

109 Parameters 

110 ---------- 

111 path 

112 Path whose file name ends in a compression suffix (see 

113 `_path_is_compressed`); convertible to 

114 `lsst.resources.ResourcePath`. 

115 

116 Returns 

117 ------- 

118 `typing.IO` [ `bytes` ] 

119 Open, seekable binary handle positioned at the start of the 

120 decompressed data. The file is deleted when the handle is 

121 closed; the caller owns closing it. 

122 

123 Raises 

124 ------ 

125 ValueError 

126 If the path is zstd-compressed and no zstd decompressor is 

127 available. 

128 """ 

129 uri = ResourcePath(path) 

130 with uri.open("rb") as raw: 

131 if uri.basename().endswith(".gz"): 

132 decompressor = cast(IO[bytes], gzip.GzipFile(fileobj=cast(IO[bytes], raw))) 

133 else: 

134 decompressor = _open_zstd_stream(cast(IO[bytes], raw)) 

135 with decompressor, contextlib.ExitStack() as stack: 

136 temp = stack.enter_context(tempfile.TemporaryFile()) 

137 shutil.copyfileobj(decompressor, temp, 1024 * 1024) 

138 temp.seek(0) 

139 # Success: ownership of the temporary file transfers to the 

140 # caller; the stack no longer closes it. 

141 stack.pop_all() 

142 return temp 

143 

144 

145@dataclasses.dataclass(frozen=True) 

146class Backend: 

147 """A file-format backend resolved from a path suffix. 

148 

149 Bundles the backend's free ``write`` function and its `InputArchive` 

150 subclass. Reading goes through the generic ``open`` / ``read`` in 

151 `lsst.images.serialization`, which use the `InputArchive`'s 

152 ``get_basic_info`` and ``open_tree``. 

153 """ 

154 

155 name: str 

156 write: Callable[..., object] 

157 input_archive: type[InputArchive] 

158 

159 

160def _backend_for_format(name: str) -> Backend | None: 

161 """Return the `Backend` registered under ``name``, or `None`. 

162 

163 Backends are imported lazily so optional dependencies (e.g. ``h5py``) 

164 are only required when actually used. 

165 """ 

166 match name: 

167 case "fits": 

168 from ..fits import FitsInputArchive 

169 from ..fits import write as fits_write 

170 

171 return Backend("fits", fits_write, FitsInputArchive) 

172 case "ndf": 

173 from ..ndf import NdfInputArchive 

174 from ..ndf import write as ndf_write 

175 

176 return Backend("ndf", ndf_write, NdfInputArchive) 

177 case "json": 

178 from ..json import JsonInputArchive 

179 from ..json import write as json_write 

180 

181 return Backend("json", json_write, JsonInputArchive) 

182 return None 

183 

184 

185def backend_for_name(name: str) -> Backend: 

186 """Return the `Backend` with the given format name. 

187 

188 Parameters 

189 ---------- 

190 name 

191 Backend format name: ``"fits"``, ``"ndf"``, or ``"json"``. 

192 

193 Raises 

194 ------ 

195 ValueError 

196 If ``name`` is not a recognized backend name. 

197 """ 

198 backend = _backend_for_format(name) 

199 if backend is None: 

200 raise ValueError(f"Unrecognized format name: {name!r}; expected one of 'fits', 'ndf', 'json'.") 

201 return backend 

202 

203 

204def backend_for_path(path: ResourcePathExpression) -> Backend: 

205 """Return the `Backend` for ``path`` based on its file extension. 

206 

207 Supported extensions: ``.fits`` (FITS), ``.h5`` / ``.sdf`` (NDF), and 

208 ``.json`` (JSON), each optionally followed by a ``.gz`` or ``.zst`` 

209 compression suffix. The NDF and FITS backends are imported lazily so 

210 optional dependencies (e.g. ``h5py``) are only required when actually 

211 used. 

212 

213 Parameters 

214 ---------- 

215 path 

216 Path whose file extension selects the backend. 

217 

218 Raises 

219 ------ 

220 ValueError 

221 If the extension is not recognized. 

222 """ 

223 uri = ResourcePath(path) 

224 name = uri.basename() 

225 for suffix in _COMPRESSION_SUFFIXES: 

226 if name.endswith(suffix): 

227 name = name.removesuffix(suffix) 

228 break 

229 match os.path.splitext(name)[1]: 

230 case ".fits": 

231 return backend_for_name("fits") 

232 case ".h5" | ".sdf": 

233 return backend_for_name("ndf") 

234 case ".json": 

235 return backend_for_name("json") 

236 case ext: 

237 raise ValueError( 

238 f"Unrecognized file extension: {ext!r} from {uri!r}; " 

239 "expected one of .fits, .h5, .sdf, .json, optionally with " 

240 "a .gz or .zst compression suffix." 

241 ) 

242 

243 

244def backend_for_stream(stream: IO[bytes]) -> Backend: 

245 """Return the `Backend` for ``stream`` based on its leading bytes. 

246 

247 The stream is restored to the position it was passed in with. 

248 Compressed content is not accepted: whoever produced the stream knows 

249 how it was compressed and decompresses it before handing it over. 

250 

251 Parameters 

252 ---------- 

253 stream 

254 Seekable binary stream positioned at the start of the data. 

255 

256 Raises 

257 ------ 

258 ValueError 

259 If the leading bytes match no supported format, including when 

260 they carry gzip or zstd compression magic. 

261 """ 

262 start = stream.tell() 

263 head = stream.read(512) 

264 stream.seek(start) 

265 if head.startswith(_FITS_MAGIC): 

266 return backend_for_name("fits") 

267 if head.startswith(_HDF5_MAGIC): 

268 return backend_for_name("ndf") 

269 if head.lstrip().startswith(b"{"): 

270 return backend_for_name("json") 

271 if head.startswith(_GZIP_MAGIC): 

272 raise ValueError("The stream appears to be gzip-compressed; decompress it before reading.") 

273 if head.startswith(_ZSTD_MAGIC): 

274 raise ValueError("The stream appears to be zstd-compressed; decompress it before reading.") 

275 raise ValueError( 

276 f"Could not identify a supported format from the leading bytes " 

277 f"{head[:16]!r}; expected FITS, HDF5/NDF, or JSON content. " 

278 "Specify the format explicitly if it is known." 

279 )