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

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 

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

13 

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""" 

20 

21from __future__ import annotations 

22 

23__all__ = ("GenericFormatter",) 

24 

25import hashlib 

26import json as _stdlib_json # disambiguates from .json subpackage 

27from collections.abc import Mapping 

28from typing import Any, ClassVar 

29 

30import astropy.io.fits 

31 

32from lsst.daf.butler import DatasetProvenance, FormatterV2 

33from lsst.resources import ResourcePath 

34 

35from . import fits as _fits 

36from . import serialization as ser 

37from .serialization import ButlerInfo, write 

38 

39 

40class GenericFormatter(FormatterV2): 

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

42 

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

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

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

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

47 

48 Notes 

49 ----- 

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

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

52 the underlying ``read`` function. 

53 """ 

54 

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

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

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

58 can_read_from_uri: ClassVar[bool] = True 

59 

60 butler_provenance: DatasetProvenance | None = None 

61 

62 # --- Write parameter handling ------------------------------------------- 

63 

64 def get_write_extension(self) -> str: 

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

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

67 ext = "." + fmt 

68 if ext not in self.supported_extensions: 

69 raise RuntimeError( 

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

71 ) 

72 return ext 

73 

74 def _validate_write_parameters(self) -> None: 

75 ext = self.get_write_extension() 

76 if ext != ".fits" and "recipe" in self.write_parameters: 

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

78 

79 @classmethod 

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

81 if not recipes: 

82 return recipes 

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

84 try: 

85 _fits.FitsCompressionOptions.model_validate(recipe) 

86 except Exception as err: 

87 err.add_note(name) 

88 raise 

89 return recipes 

90 

91 # --- Write path --------------------------------------------------------- 

92 

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

94 self._validate_write_parameters() 

95 ext = self.get_write_extension() 

96 butler_info = ButlerInfo( 

97 dataset=self.dataset_ref.to_simple(), 

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

99 ) 

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

101 if ext == ".fits": 

102 kwargs["update_header"] = self._update_header 

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

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

105 # The generic write() dispatches to the FITS / JSON / NDF backend by 

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

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

108 

109 def add_provenance( 

110 self, 

111 in_memory_dataset: Any, 

112 /, 

113 *, 

114 provenance: DatasetProvenance | None = None, 

115 ) -> Any: 

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

117 # rather than mutating the dataset. 

118 self.butler_provenance = provenance 

119 return in_memory_dataset 

120 

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

122 

123 def _get_compression_seed(self) -> int: 

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

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

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

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

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

129 # simpler. 

130 hash_bytes = hashlib.md5( 

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

132 usedforsecurity=False, 

133 ).digest() 

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

135 # 10000] range allowed by FITS. 

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

137 

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

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

140 try: 

141 config = self.write_recipes[recipe] 

142 except KeyError: 

143 if recipe == "default": 

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

145 return {} 

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

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

148 

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

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

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

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

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

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

155 for key in list(header): 

156 if key.startswith("LSST BUTLER"): 

157 del header[key] 

158 if self.butler_provenance is not None: 

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

160 self.dataset_ref, 

161 prefix="HIERARCH LSST BUTLER", 

162 sep=" ", 

163 simple_types=True, 

164 max_inputs=3_000, 

165 ).items(): 

166 header.set(key, value) 

167 

168 # --- Read path --------------------------------------------------------- 

169 

170 def read_from_uri( 

171 self, 

172 uri: ResourcePath, 

173 component: str | None = None, 

174 expected_size: int = -1, 

175 ) -> Any: 

176 kwargs = self.file_descriptor.parameters or {} 

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

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

179 if component is None: 

180 return reader.read(**kwargs) 

181 return reader.get_component(component, **kwargs)