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

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() dispatches to the FITS / JSON / NDF backend by

129 # the file extension, which get_write_extension has already set on uri.

130 write(in_memory_dataset, uri.ospath, **kwargs)

131

132 def add_provenance(

133 self,

134 in_memory_dataset: Any,

135 /,

136 *,

137 provenance: DatasetProvenance | None = None,

138 ) -> Any:

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

140 # rather than mutating the dataset.

141 self.butler_provenance = provenance

142 return in_memory_dataset

143

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

145

146 def _get_compression_seed(self) -> int:

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

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

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

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

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

152 # simpler.

153 hash_bytes = hashlib.md5(

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

155 usedforsecurity=False,

156 ).digest()

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

158 # 10000] range allowed by FITS.

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

160

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

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

163 try:

164 config = self.write_recipes[recipe]

165 except KeyError:

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

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

168 return {}

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

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

171

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

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

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

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

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

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

178 for key in list(header):

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

180 del header[key]

181 if self.butler_provenance is not None:

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

183 self.dataset_ref,

184 prefix="HIERARCH LSST BUTLER",

185 sep=" ",

186 simple_types=True,

187 max_inputs=3_000,

188 ).items():

189 header.set(key, value)

190

191 # --- Component tree cache -----------------------------------------------

192

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

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

195

196 Parameters

197 ----------

198 component

199 Name of the component to read.

200

201 Returns

202 -------

203 hit : `bool`

204 Whether the component could be served from the cache.

205 value

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

207

208 Raises

209 ------

210 lsst.images.serialization.InvalidComponentError

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

212 """

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

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

215 return False, None

216 try:

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

218 except ser.ArchiveAccessRequiredError:

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

220 # the file has to be opened and read.

221 return False, None

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

223

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

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

226 same dataset can be served without reopening the file.

227 """

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

229

230 @staticmethod

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

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

233

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

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

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

237 each other through the cache.

238 """

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

240 return copy.deepcopy(value)

241 return value

242

243 # --- Read path ---------------------------------------------------------

244

245 def read_from_uri(

246 self,

247 uri: ResourcePath,

248 component: str | None = None,

249 expected_size: int = -1,

250 ) -> Any:

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

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

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

254 # than the full file.

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

256 return NotImplemented

257

258 # Now call the generalized reader.

259 return self._read_from_resource_path(uri, component)

260

261 def _resolve_multi_component_request(

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

263 ) -> set[str]:

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

265

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

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

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

269

270 Raises

271 ------

272 RuntimeError

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

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

275 """

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

277 if requested_components is None:

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

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

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

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

282 # for similar cases in the future.

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

284 if not requested_components:

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

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

287 components = set(ensure_iterable(requested_components))

288 if "components" in components:

289 raise RuntimeError(

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

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

292 )

293 return components

294

295 def _build_component_kwargs(

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

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

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

299 understands.

300

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

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

303 declare it.

304

305 Raises

306 ------

307 RuntimeError

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

309 if a supplied parameter is not understood by any requested

310 component.

311 """

312 used_parameter_keys = set()

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

314 for comp in components:

315 if comp not in all_components:

316 raise RuntimeError(

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

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

319 )

320 component_kwargs[comp] = {}

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

322 if param in all_components[comp].parameters:

323 component_kwargs[comp][param] = value

324 used_parameter_keys.add(param)

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

326 raise RuntimeError(

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

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

329 )

330 return component_kwargs

331

332 def _read_components(

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

334 ) -> dict[str, Any]:

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

336

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

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

339 """

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

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

342 if not params:

343 hit, value = self._component_from_cache(comp)

344 if hit:

345 components_to_return[comp] = value

346

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

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

349 # to read the rest.

350 with ser.open(uri, cls=pytype, partial=True) as reader:

351 tree = reader.get_tree()

352 self._cache_tree(tree)

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

354 if comp not in components_to_return:

355 components_to_return[comp] = self._detach_component(

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

357 )

358 return components_to_return

359

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

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

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

363 # to be made regarding caching.

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

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

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

367

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

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

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

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

372 return NotImplemented

373

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

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

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

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

378 components: set[str] = set()

379 want_component_dict = False

380 if component == "components":

381 want_component_dict = True

382 components = self._resolve_multi_component_request(kwargs, all_components)

383 elif component:

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

385 components = {component}

386

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

388 raise RuntimeError(

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

390 )

391

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

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

394 if component is None:

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

396 tree = reader.get_tree()

397 self._cache_tree(tree)

398 # Cutout read.

399 return reader.read(**kwargs)

400

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

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

403

404 if want_component_dict:

405 return components_to_return

406 return components_to_return.popitem()[1]

407

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

409 # Docstring inherited.

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

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

412 # can trigger a cache load.

413 return self._read_from_resource_path(path, component)