Coverage for python/lsst/images/json/_output_archive.py: 59%

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 

12from __future__ import annotations 

13 

14__all__ = ("JsonOutputArchive", "write") 

15 

16from collections.abc import Callable, Hashable, Iterator, Mapping 

17from typing import TYPE_CHECKING, Any 

18 

19import astropy.table 

20import numpy as np 

21import pydantic 

22 

23from lsst.resources import ResourcePath, ResourcePathExpression 

24 

25from .._transforms import FrameSet 

26from ..serialization import ( 

27 ArchiveTree, 

28 ButlerInfo, 

29 InlineArrayModel, 

30 JsonRef, 

31 MetadataValue, 

32 NestedOutputArchive, 

33 NumberType, 

34 OutputArchive, 

35 TableColumnModel, 

36 TableModel, 

37 no_header_updates, 

38) 

39 

40if TYPE_CHECKING: 

41 import astropy.io.fits 

42 

43 

44def write( 

45 obj: Any, 

46 path: ResourcePathExpression | None = None, 

47 metadata: dict[str, MetadataValue] | None = None, 

48 butler_info: ButlerInfo | None = None, 

49) -> ArchiveTree: 

50 """Write an object with a ``serialize`` method to a JSON file. 

51 

52 Parameters 

53 ---------- 

54 path 

55 Name of the file to write to (may be a URI). If not provided, a 

56 serializable model is returned but not written to disk. 

57 metadata 

58 Additional metadata to save with the object. This will override any 

59 flexible metadata carried by the object itself with the same keys. 

60 butler_info 

61 Butler information to store in the file. 

62 

63 Returns 

64 ------- 

65 `.serialization.ArchiveTree` 

66 The serialized representation of the object. 

67 """ 

68 archive = JsonOutputArchive() 

69 name = getattr(obj, "_archive_default_name", None) 

70 tree = archive.serialize_direct(name, obj.serialize) if name is not None else obj.serialize(archive) 

71 if metadata is not None: 71 ↛ 72line 71 didn't jump to line 72 because the condition on line 71 was never true

72 tree.metadata.update(metadata) 

73 if butler_info is not None: 

74 tree.butler_info = butler_info 

75 archive.finish(tree) 

76 if path is not None: 76 ↛ 78line 76 didn't jump to line 78 because the condition on line 76 was always true

77 ResourcePath(path).write(tree.model_dump_json().encode()) 

78 return tree 

79 

80 

81class JsonOutputArchive(OutputArchive[JsonRef]): 

82 """An implementation of the `.serialization.OutputArchive` interface that 

83 writes to JSON files. 

84 

85 This archive type is designed for pure-JSON objects and cases where any 

86 images or tables are tiny. It will be *extremely* inefficient for large 

87 images or tables, if it works at all. 

88 """ 

89 

90 def __init__(self) -> None: 

91 super().__init__() 

92 self._pointers: dict[Hashable, JsonRef] = {} 

93 self._indirect: list[Any] = [] 

94 self._frame_sets: list[tuple[FrameSet, JsonRef]] = [] 

95 

96 def serialize_direct[T: pydantic.BaseModel | None]( 

97 self, name: str, serializer: Callable[[OutputArchive[JsonRef]], T] 

98 ) -> T: 

99 nested = NestedOutputArchive[JsonRef](name, self) 

100 return serializer(nested) 

101 

102 def serialize_pointer[T: ArchiveTree]( 

103 self, name: str, serializer: Callable[[OutputArchive[JsonRef]], T], key: Hashable 

104 ) -> JsonRef: 

105 if (pointer := self._pointers.get(key)) is not None: 

106 return pointer 

107 pointer = JsonRef.model_construct(ref=f"#/indirect/{len(self._indirect)}") 

108 self._indirect.append(self.serialize_direct(name, serializer).model_dump()) 

109 self._pointers[key] = pointer 

110 return pointer 

111 

112 def serialize_frame_set[T: ArchiveTree]( 

113 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable 

114 ) -> JsonRef: 

115 pointer = self.serialize_pointer(name, serializer, key) 

116 self._frame_sets.append((frame_set, pointer)) 

117 return pointer 

118 

119 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, JsonRef]]: 

120 return iter(self._frame_sets) 

121 

122 def add_array( 

123 self, 

124 array: np.ndarray, 

125 *, 

126 name: str | None = None, 

127 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

128 tile_shape: tuple[int, ...] | None = None, 

129 options_name: str | None = None, 

130 ) -> InlineArrayModel: 

131 return InlineArrayModel( 

132 data=array.tolist(), 

133 datatype=NumberType.from_numpy(array.dtype), 

134 ) 

135 

136 def add_table( 

137 self, 

138 table: astropy.table.Table, 

139 *, 

140 name: str | None = None, 

141 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

142 ) -> TableModel: 

143 columns = TableColumnModel.from_table(table, inline=True) 

144 return TableModel(columns=columns, meta=table.meta) 

145 

146 def add_structured_array( 

147 self, 

148 array: np.ndarray, 

149 *, 

150 name: str | None = None, 

151 units: Mapping[str, astropy.units.Unit] | None = None, 

152 descriptions: Mapping[str, str] | None = None, 

153 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

154 ) -> TableModel: 

155 columns = TableColumnModel.from_record_array(array, inline=True) 

156 for c in columns: 

157 if units and (unit := units.get(c.name)): 

158 c.unit = unit 

159 if descriptions and (description := descriptions.get(c.name)): 

160 c.description = description 

161 return TableModel(columns=columns) 

162 

163 def finish[T: ArchiveTree](self, tree: T) -> T: 

164 """Finish serialization. 

165 

166 Parameters 

167 ---------- 

168 tree 

169 Serialized archive tree to write, which is modified in place 

170 (the ``indirect`` attribute is overwritten) and then returned. 

171 """ 

172 tree.indirect = self._indirect 

173 return tree