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

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 ._geom import Box

42from .serialization import ButlerInfo, write_archive

45class _TreeCache(NamedTuple):

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

47 serialization tree.

48 """

50 id_: uuid.UUID | None = None

51 tree: ser.ArchiveTree | None = None

54_DETACHED_ARCHIVE = ser.DetachedArchive()

57class GenericFormatter(FormatterV2):

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

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

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

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

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

65 Notes

66 -----

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

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

69 the underlying ``read`` function.

70 """

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

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

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

75 can_read_from_uri: ClassVar[bool] = True

76 can_read_from_local_file: ClassVar[bool] = True

78 butler_provenance: DatasetProvenance | None = None

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

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

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

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

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

86 # --- Write parameter handling -------------------------------------------

88 def get_write_extension(self) -> str:

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

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

91 ext = "." + fmt

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

93 raise RuntimeError(

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

95 )

96 return ext

98 def _validate_write_parameters(self) -> None:

99 ext = self.get_write_extension()

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

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

102

103 @classmethod

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

105 if not recipes:

106 return recipes

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

108 try:

109 _fits.FitsCompressionOptions.model_validate(recipe)

110 except Exception as err:

111 err.add_note(name)

112 raise

113 return recipes

114

115 # --- Write path ---------------------------------------------------------

116

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

118 self._validate_write_parameters()

119 ext = self.get_write_extension()

120 butler_info = ButlerInfo(

121 dataset=self.dataset_ref.to_simple(),

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

123 )

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

125 if ext == ".fits":

126 kwargs["update_header"] = self._update_header

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

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

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

130 # backend by the file extension, which get_write_extension has

131 # already set on uri.

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

133

134 def add_provenance(

135 self,

136 in_memory_dataset: Any,

137 /,

138 *,

139 provenance: DatasetProvenance | None = None,

140 ) -> Any:

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

142 # rather than mutating the dataset.

143 self.butler_provenance = provenance

144 return in_memory_dataset

145

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

147

148 def _get_compression_seed(self) -> int:

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

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

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

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

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

154 # simpler.

155 hash_bytes = hashlib.md5(

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

157 usedforsecurity=False,

158 ).digest()

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

160 # 10000] range allowed by FITS.

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

162

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

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

165 try:

166 config = self.write_recipes[recipe]

167 except KeyError:

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

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

170 return {}

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

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

173

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

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

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

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

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

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

180 for key in list(header):

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

182 del header[key]

183 if self.butler_provenance is not None:

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

185 self.dataset_ref,

186 prefix="HIERARCH LSST BUTLER",

187 sep=" ",

188 simple_types=True,

189 max_inputs=3_000,

190 ).items():

191 header.set(key, value)

192

193 # --- Component tree cache -----------------------------------------------

194

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

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

197

198 Parameters

199 ----------

200 component

201 Name of the component to read.

202

203 Returns

204 -------

205 hit : `bool`

206 Whether the component could be served from the cache.

207 value

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

209

210 Raises

211 ------

212 lsst.images.serialization.InvalidComponentError

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

214 """

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

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

217 return False, None

218 try:

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

220 except ser.ArchiveAccessRequiredError:

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

222 # the file has to be opened and read.

223 return False, None

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

225

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

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

228 same dataset can be served without reopening the file.

229 """

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

231

232 @staticmethod

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

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

235

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

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

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

239 each other through the cache.

240 """

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

242 return copy.deepcopy(value)

243 return value

244

245 # --- Read path ---------------------------------------------------------

246

247 def read_from_uri(

248 self,

249 uri: ResourcePath,

250 component: str | None = None,

251 expected_size: int = -1,

252 ) -> Any:

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

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

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

256 # than the full file.

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

258 return NotImplemented

259

260 # Now call the generalized reader.

261 return self._read_from_resource_path(uri, component)

262

263 def _resolve_multi_component_request(

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

265 ) -> set[str]:

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

267

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

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

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

271

272 Raises

273 ------

274 RuntimeError

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

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

277 """

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

279 if requested_components is None:

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

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

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

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

284 # for similar cases in the future.

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

286 if not requested_components:

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

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

289 components = set(ensure_iterable(requested_components))

290 if "components" in components:

291 raise RuntimeError(

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

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

294 )

295 return components

296

297 def _build_component_kwargs(

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

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

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

301 understands.

302

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

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

305 declare it.

306

307 Raises

308 ------

309 RuntimeError

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

311 if a supplied parameter is not understood by any requested

312 component.

313 """

314 used_parameter_keys = set()

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

316 for comp in components:

317 if comp not in all_components:

318 raise RuntimeError(

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

320 f"by storage class {self.file_descriptor.storageClass.name}."

321 )

322 component_kwargs[comp] = {}

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

324 if param in all_components[comp].parameters:

325 component_kwargs[comp][param] = value

326 used_parameter_keys.add(param)

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

328 raise RuntimeError(

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

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

331 )

332 return component_kwargs

333

334 def _read_components(

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

336 ) -> dict[str, Any]:

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

338

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

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

341 """

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

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

344 if not params:

345 hit, value = self._component_from_cache(comp)

346 if hit:

347 components_to_return[comp] = value

348

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

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

351 # to read the rest.

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

353 tree = reader.get_tree()

354 self._cache_tree(tree)

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

356 if comp not in components_to_return:

357 components_to_return[comp] = self._detach_component(

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

359 )

360 return components_to_return

361

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

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

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

365 # to be made regarding caching.

366 kwargs = self._coerce_legacy_bboxes(self.file_descriptor.parameters or {})

367 write_storage_class = self.file_descriptor.storageClass

368 pytype: type[Any] = write_storage_class.pytype

369 all_components = write_storage_class.allComponents()

370

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

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

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

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

375 return NotImplemented

376

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

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

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

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

381 components: set[str] = set()

382 want_component_dict = False

383 if component == "components":

384 want_component_dict = True

385 components = self._resolve_multi_component_request(kwargs, all_components)

386 elif component:

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

388 components = {component}

389

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

391 raise RuntimeError(

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

393 )

394

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

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

397 if component is None:

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

399 tree = reader.get_tree()

400 self._cache_tree(tree)

401 # Cutout read.

402 return reader.read(**kwargs)

403

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

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

406

407 if want_component_dict:

408 return components_to_return

409 return components_to_return.popitem()[1]

410

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

412 # Docstring inherited.

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

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

415 # can trigger a cache load.

416 return self._read_from_resource_path(path, component)

417

418 @staticmethod

419 def _coerce_legacy_bboxes(parameters: dict[str, Any]) -> dict[str, Any]:

420 """Modify a parameters dictionary in place by calling

421 `.Box.from_legacy` on the values associated with any 'bbox' keys that

422 are not already `.Box` instances.

423 """

424 for k, v in parameters.items():

425 if k == "bbox" and not isinstance(v, Box): 425 ↛ 426line 425 didn't jump to line 426 because the condition on line 425 was never true

426 v = Box.from_legacy(v)

427 parameters[k] = v

428 return parameters