Coverage for python/lsst/images/formatters.py: 92%

1# This file is part of lsst-images.

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.

9# Use of this source code is governed by a 3-clause BSD-style

10# license that can be found in the LICENSE file.

12"""Unified butler formatter for lsst.images.

14This formatter dispatches on a write-time ``format`` parameter and on the

15file extension at read time, replacing the three per-format

16(`lsst.images.fits.formatters`, `lsst.images.json.formatters`,

17`lsst.images.ndf.formatters`) hierarchies that previously duplicated almost

18all of their logic.

19"""

21from __future__ import annotations

23__all__ = ("GenericFormatter",)

25import copy

26import hashlib

27import json as _stdlib_json # disambiguates from .json subpackage

28import threading

29import uuid

30from collections.abc import Mapping

31from typing import Any, ClassVar, NamedTuple

33import astropy.io.fits

35from lsst.daf.butler import DatasetProvenance, FormatterV2

36from lsst.resources import ResourcePath, ResourcePathExpression

37from lsst.utils.iteration import ensure_iterable

39from . import fits as _fits

40from . import serialization as ser

41from .serialization import ButlerInfo, write_archive

44class _TreeCache(NamedTuple):

45 """Single-slot cache pairing a dataset ID with its validated

46 serialization tree.

47 """

49 id_: uuid.UUID | None = None

50 tree: ser.ArchiveTree | None = None

53_DETACHED_ARCHIVE = ser.DetachedArchive()

56class GenericFormatter(FormatterV2):

57 """Unified butler formatter for any lsst.images type.

59 The on-disk format is selected by the ``format`` write parameter

60 (``fits``, ``json``, ``sdf``) at write time and by the file

61 extension at read time. The default format is taken from

62 ``self.default_extension`` (``.fits`` for the base class).

64 Notes

65 -----

66 Subclasses (`ImageFormatter` and below) add component-level read

67 support. This base class forwards any read parameters straight to

68 the underlying ``read`` function.

69 """

71 default_extension: ClassVar[str] = ".fits"

72 supported_extensions: ClassVar[frozenset[str]] = frozenset({".fits", ".sdf", ".json"})

73 supported_write_parameters: ClassVar[frozenset[str]] = frozenset({"format", "recipe"})

74 can_read_from_uri: ClassVar[bool] = True

75 can_read_from_local_file: ClassVar[bool] = True

77 butler_provenance: DatasetProvenance | None = None

79 # Most recently read serialization tree, kept so that repeated component

80 # reads of the same dataset do not reopen the file. It is thread-local:

81 # each thread keeps its own most-recent tree, so concurrent reads in

82 # different threads never invalidate each other and no locking is needed.

83 _tree_cache: ClassVar[threading.local] = threading.local()

85 # --- Write parameter handling -------------------------------------------

87 def get_write_extension(self) -> str:

88 default_fmt = self.default_extension.lstrip(".")

89 fmt = self.write_parameters.get("format", default_fmt)

90 ext = "." + fmt

91 if ext not in self.supported_extensions: 91 ↛ 92line 91 didn't jump to line 92 because the condition on line 91 was never true

92 raise RuntimeError(

93 f"Requested format {fmt!r} is not supported; expected one of {{fits, json, sdf}}."

94 )

95 return ext

97 def _validate_write_parameters(self) -> None:

98 ext = self.get_write_extension()

99 if ext != ".fits" and "recipe" in self.write_parameters: 99 ↛ 100line 99 didn't jump to line 100 because the condition on line 99 was never true

100 raise RuntimeError("The 'recipe' write parameter is only valid for FITS output.")

101

102 @classmethod

103 def validate_write_recipes(cls, recipes: Mapping[str, Any] | None) -> Mapping[str, Any] | None:

104 if not recipes:

105 return recipes

106 for name, recipe in recipes.items():

107 try:

108 _fits.FitsCompressionOptions.model_validate(recipe)

109 except Exception as err:

110 err.add_note(name)

111 raise

112 return recipes

113

114 # --- Write path ---------------------------------------------------------

115

116 def write_local_file(self, in_memory_dataset: Any, uri: ResourcePath) -> None:

117 self._validate_write_parameters()

118 ext = self.get_write_extension()

119 butler_info = ButlerInfo(

120 dataset=self.dataset_ref.to_simple(),

121 provenance=self.butler_provenance if self.butler_provenance is not None else DatasetProvenance(),

122 )

123 kwargs: dict[str, Any] = {"butler_info": butler_info}

124 if ext == ".fits":

125 kwargs["update_header"] = self._update_header

126 kwargs["compression_options"] = self._get_compression_options()

127 kwargs["compression_seed"] = self._get_compression_seed()

128 # The generic write_archive() dispatches to the FITS / JSON / NDF

129 # backend by the file extension, which get_write_extension has

130 # already set on uri.

131 write_archive(in_memory_dataset, uri.ospath, **kwargs)

132

133 def add_provenance(

134 self,

135 in_memory_dataset: Any,

136 /,

137 *,

138 provenance: DatasetProvenance | None = None,

139 ) -> Any:

140 # A FormatterV2 instance is used once; stash provenance on self

141 # rather than mutating the dataset.

142 self.butler_provenance = provenance

143 return in_memory_dataset

144

145 # --- FITS-specific helpers (kept verbatim from fits/formatters.py) ----

146

147 def _get_compression_seed(self) -> int:

148 # Set the seed based on data ID (all logic here duplicated from

149 # obs_base). We can't just use 'hash', since like 'set' that's not

150 # deterministic. And we can't rely on a DimensionPacker because those

151 # are only defined for certain combinations of dimensions. Doing an MD5

152 # of the JSON feels like overkill but I don't really see anything much

153 # simpler.

154 hash_bytes = hashlib.md5(

155 _stdlib_json.dumps(list(self.data_id.required_values)).encode(),

156 usedforsecurity=False,

157 ).digest()

158 # And it *really* feels like overkill when we squash that into the [1,

159 # 10000] range allowed by FITS.

160 return 1 + int.from_bytes(hash_bytes) % 9999

161

162 def _get_compression_options(self) -> dict[str, _fits.FitsCompressionOptions]:

163 recipe = self.write_parameters.get("recipe", "default")

164 try:

165 config = self.write_recipes[recipe]

166 except KeyError:

167 if recipe == "default": 167 ↛ 170line 167 didn't jump to line 170 because the condition on line 167 was always true

168 # If there's no default recipe just use the software defaults.

169 return {}

170 raise RuntimeError(f"Invalid recipe for GenericFormatter: {recipe!r}.") from None

171 return {k: _fits.FitsCompressionOptions.model_validate(v) for k, v in config.items()}

172

173 def _update_header(self, header: astropy.io.fits.Header) -> None:

174 # Logic here largely lifted from lsst.obs.base.utils, which we

175 # can't use directly for dependency and maybe mapping-type

176 # (PropertyList vs. astropy) reasons. We assume we can always add

177 # long cards (astropy will CONTINUE them) but not comments

178 # (astropy will truncate and warn on long cards).

179 for key in list(header):

180 if key.startswith("LSST BUTLER"): 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true

181 del header[key]

182 if self.butler_provenance is not None:

183 for key, value in self.butler_provenance.to_flat_dict(

184 self.dataset_ref,

185 prefix="HIERARCH LSST BUTLER",

186 sep=" ",

187 simple_types=True,

188 max_inputs=3_000,

189 ).items():

190 header.set(key, value)

191

192 # --- Component tree cache -----------------------------------------------

193

194 def _component_from_cache(self, component: str) -> tuple[bool, Any]:

195 """Try to deserialize a component from the most recently read tree.

196

197 Parameters

198 ----------

199 component

200 Name of the component to read.

201

202 Returns

203 -------

204 hit : `bool`

205 Whether the component could be served from the cache.

206 value

207 The deserialized component; `None` on a cache miss.

208

209 Raises

210 ------

211 lsst.images.serialization.InvalidComponentError

212 Raised if the cached tree does not recognize ``component``.

213 """

214 cache = getattr(type(self)._tree_cache, "value", _TreeCache())

215 if cache.tree is None or cache.id_ != self.dataset_ref.id:

216 return False, None

217 try:

218 value = cache.tree.deserialize_component(component, _DETACHED_ARCHIVE)

219 except ser.ArchiveAccessRequiredError:

220 # The component points at data stored outside the JSON tree, so

221 # the file has to be opened and read.

222 return False, None

223 return True, self._detach_component(cache.tree, component, value)

224

225 def _cache_tree(self, tree: ser.ArchiveTree) -> None:

226 """Remember a validated tree so that later component reads of the

227 same dataset can be served without reopening the file.

228 """

229 type(self)._tree_cache.value = _TreeCache(id_=self.dataset_ref.id, tree=tree)

230

231 @staticmethod

232 def _detach_component(tree: ser.ArchiveTree, component: str, value: Any) -> Any:

233 """Copy a component value if it is owned by the given tree.

234

235 `~lsst.images.serialization.ArchiveTree.deserialize_component`

236 returns plain (non-tree) models by reference from the tree, so

237 without a copy repeated reads of a mutable component would alias

238 each other through the cache.

239 """

240 if value is not None and value is getattr(tree, component, None):

241 return copy.deepcopy(value)

242 return value

243

244 # --- Read path ---------------------------------------------------------

245

246 def read_from_uri(

247 self,

248 uri: ResourcePath,

249 component: str | None = None,

250 expected_size: int = -1,

251 ) -> Any:

252 # For full read, always use local file read since the entire file has

253 # to be read anyhow and we should allow it to be cached. Cutouts

254 # can use remote reads since that is generally less to be downloaded

255 # than the full file.

256 if not component and not self.file_descriptor.parameters: 256 ↛ 257line 256 didn't jump to line 257 because the condition on line 256 was never true

257 return NotImplemented

258

259 # Now call the generalized reader.

260 return self._read_from_resource_path(uri, component)

261

262 def _resolve_multi_component_request(

263 self, kwargs: dict[str, Any], all_components: Mapping[str, Any]

264 ) -> set[str]:

265 """Resolve the component set for a ``"components"`` multi-read request.

266

267 Consumes the ``components`` entry from ``kwargs`` (the list of

268 requested component names) and returns the set of component names to

269 read. If no explicit list was given, every component is returned.

270

271 Raises

272 ------

273 RuntimeError

274 Raised if an explicit but empty request is given, or if the

275 ``"components"`` pseudo-component is itself listed.

276 """

277 requested_components = kwargs.pop("components", None)

278 if requested_components is None:

279 # No explicit request, so read every component. Drop the

280 # "components" pseudo-component itself and masked_image (its pixels

281 # are already covered by the individual image, mask, and variance

282 # components). Hard-coding this is not ideal so we have to watch

283 # for similar cases in the future.

284 return {c for c in all_components if c not in {"components", "masked_image"}}

285 if not requested_components:

286 raise RuntimeError("Requesting multiple components but received empty request.")

287 # Force to a set in case someone has tried doing "components=x".

288 components = set(ensure_iterable(requested_components))

289 if "components" in components:

290 raise RuntimeError(

291 "The 'components' component should not be specified in the 'components' parameter. "

292 "To request all components, do not specify any value for the 'components' parameter."

293 )

294 return components

295

296 def _build_component_kwargs(

297 self, components: set[str], kwargs: dict[str, Any], all_components: Mapping[str, Any]

298 ) -> dict[str, dict[str, Any]]:

299 """Map each requested component to the subset of parameters it

300 understands.

301

302 This lets a single read ask for, for example, an image cutout

303 alongside a PSF: each parameter is routed to the components that

304 declare it.

305

306 Raises

307 ------

308 RuntimeError

309 Raised if a requested component is unknown to the storage class, or

310 if a supplied parameter is not understood by any requested

311 component.

312 """

313 used_parameter_keys = set()

314 component_kwargs: dict[str, dict[str, Any]] = {}

315 for comp in components:

316 if comp not in all_components:

317 raise RuntimeError(

318 f"Requested data for component {comp} but that component is not understood "

319 f"by storage class {self.dataset_ref.datasetType.storageClass.name}."

320 )

321 component_kwargs[comp] = {}

322 for param, value in kwargs.items():

323 if param in all_components[comp].parameters:

324 component_kwargs[comp][param] = value

325 used_parameter_keys.add(param)

326 if kwargs and (unused := (set(kwargs) - used_parameter_keys)):

327 raise RuntimeError(

328 f"Specified parameters ({unused}) that are not known to any of the "

329 f"requested components ({components})."

330 )

331 return component_kwargs

332

333 def _read_components(

334 self, uri: ResourcePathExpression, pytype: type[Any], component_kwargs: dict[str, dict[str, Any]]

335 ) -> dict[str, Any]:

336 """Read the requested components into a name-keyed dict.

337

338 Parameterless components are served from the cached tree when possible;

339 the file is opened only if some components are not cached.

340 """

341 components_to_return: dict[str, Any] = {}

342 for comp, params in component_kwargs.items():

343 if not params:

344 hit, value = self._component_from_cache(comp)

345 if hit:

346 components_to_return[comp] = value

347

348 if len(components_to_return) != len(component_kwargs):

349 # Some components were not available in the cache, so open the file

350 # to read the rest.

351 with ser.open_archive(uri, cls=pytype, partial=True) as reader:

352 tree = reader.get_tree()

353 self._cache_tree(tree)

354 for comp, params in component_kwargs.items():

355 if comp not in components_to_return:

356 components_to_return[comp] = self._detach_component(

357 tree, comp, reader.get_component(comp, **params)

358 )

359 return components_to_return

360

361 def _read_from_resource_path(self, uri: ResourcePathExpression, component: str | None = None) -> Any:

362 # General purpose reader that can be called with both local and remote

363 # file. The URI and local file readers are distinct to allow decisions

364 # to be made regarding caching.

365 kwargs = dict(self.file_descriptor.parameters or {})

366 pytype: type[Any] = self.dataset_ref.datasetType.storageClass.pytype

367 all_components = self.dataset_ref.datasetType.storageClass.allComponents()

368

369 # An all-components request with no parameters needs the whole file, so

370 # on a remote URI defer to the local-file read that can cache it (this

371 # mirrors the full-read check in read_from_uri).

372 if component == "components" and not kwargs and isinstance(uri, ResourcePath) and not uri.isLocal: 372 ↛ 373line 372 didn't jump to line 373 because the condition on line 372 was never true

373 return NotImplemented

374

375 # The "components" pseudo-component is special: rather than modifying a

376 # single component it asks for several components to be returned

377 # together as a dict. Everything else is a single named component.

378 # A set ensures we never ask for the same component twice.

379 components: set[str] = set()

380 want_component_dict = False

381 if component == "components":

382 want_component_dict = True

383 components = self._resolve_multi_component_request(kwargs, all_components)

384 elif component:

385 # Simplify the logic below so we only ever deal with a set.

386 components = {component}

387

388 if "components" in kwargs: 388 ↛ 389line 388 didn't jump to line 389 because the condition on line 388 was never true

389 raise RuntimeError(

390 "Multiple component requests can only be specified if you use the 'components' component."

391 )

392

393 # If this is not a component read but a full read with parameters,

394 # do that now before we focus on the component logic.

395 if component is None:

396 with ser.open_archive(uri, cls=pytype, partial=bool(kwargs)) as reader:

397 tree = reader.get_tree()

398 self._cache_tree(tree)

399 # Cutout read.

400 return reader.read(**kwargs)

401

402 component_kwargs = self._build_component_kwargs(components, kwargs, all_components)

403 components_to_return = self._read_components(uri, pytype, component_kwargs)

404

405 if want_component_dict:

406 return components_to_return

407 return components_to_return.popitem()[1]

408

409 def read_from_local_file(self, path: str, component: str | None = None, expected_size: int = -1) -> Any:

410 # Docstring inherited.

411 # Call the generalized reader that does not care whether this is

412 # a local or remote file. The distinction exists here to ensure we

413 # can trigger a cache load.

414 return self._read_from_resource_path(path, component)