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 typing import Any, ClassVar 

28 

29import astropy.io.fits 

30 

31from lsst.daf.butler import DatasetProvenance, FormatterV2 

32from lsst.resources import ResourcePath 

33 

34from . import fits as _fits 

35from . import serialization as ser 

36from .serialization import ButlerInfo, write 

37 

38 

39class GenericFormatter(FormatterV2): 

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

41 

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

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

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

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

46 

47 Notes 

48 ----- 

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

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

51 the underlying ``read`` function. 

52 """ 

53 

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

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

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

57 can_read_from_uri: ClassVar[bool] = True 

58 

59 butler_provenance: DatasetProvenance | None = None 

60 

61 # --- Write parameter handling ------------------------------------------- 

62 

63 def get_write_extension(self) -> str: 

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

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

66 ext = "." + fmt 

67 if ext not in self.supported_extensions: 

68 raise RuntimeError( 

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

70 ) 

71 return ext 

72 

73 def _validate_write_parameters(self) -> None: 

74 ext = self.get_write_extension() 

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

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

77 

78 # --- Write path --------------------------------------------------------- 

79 

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

81 self._validate_write_parameters() 

82 ext = self.get_write_extension() 

83 butler_info = ButlerInfo( 

84 dataset=self.dataset_ref.to_simple(), 

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

86 ) 

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

88 if ext == ".fits": 

89 kwargs["update_header"] = self._update_header 

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

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

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

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

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

95 

96 def add_provenance( 

97 self, 

98 in_memory_dataset: Any, 

99 /, 

100 *, 

101 provenance: DatasetProvenance | None = None, 

102 ) -> Any: 

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

104 # rather than mutating the dataset. 

105 self.butler_provenance = provenance 

106 return in_memory_dataset 

107 

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

109 

110 def _get_compression_seed(self) -> int: 

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

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

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

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

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

116 # simpler. 

117 hash_bytes = hashlib.md5( 

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

119 usedforsecurity=False, 

120 ).digest() 

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

122 # 10000] range allowed by FITS. 

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

124 

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

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

127 try: 

128 config = self.write_recipes[recipe] 

129 except KeyError: 

130 if recipe == "default": 

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

132 return {} 

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

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

135 

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

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

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

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

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

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

142 for key in list(header): 

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

144 del header[key] 

145 if self.butler_provenance is not None: 

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

147 self.dataset_ref, 

148 prefix="HIERARCH LSST BUTLER", 

149 sep=" ", 

150 simple_types=True, 

151 max_inputs=3_000, 

152 ).items(): 

153 header.set(key, value) 

154 

155 # --- Read path --------------------------------------------------------- 

156 

157 def read_from_uri( 

158 self, 

159 uri: ResourcePath, 

160 component: str | None = None, 

161 expected_size: int = -1, 

162 ) -> Any: 

163 kwargs = self.file_descriptor.parameters or {} 

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

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

166 if component is None: 

167 return reader.read(**kwargs) 

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