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

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 obj 

55 Object with a ``serialize`` method to write. 

56 path 

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

58 serializable model is returned but not written to disk. 

59 metadata 

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

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

62 butler_info 

63 Butler information to store in the file. 

64 

65 Returns 

66 ------- 

67 `.serialization.ArchiveTree` 

68 The serialized representation of the object. 

69 """ 

70 archive = JsonOutputArchive() 

71 tree = archive.serialize_root(obj, metadata, butler_info) 

72 archive.finish(tree) 

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

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

75 return tree 

76 

77 

78class JsonOutputArchive(OutputArchive[JsonRef]): 

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

80 writes to JSON files. 

81 

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

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

84 images or tables, if it works at all. 

85 """ 

86 

87 def __init__(self) -> None: 

88 super().__init__() 

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

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

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

92 

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

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

95 ) -> T: 

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

97 return serializer(nested) 

98 

99 def serialize_pointer[T: ArchiveTree]( 

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

101 ) -> JsonRef: 

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

103 return pointer 

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

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

106 self._pointers[key] = pointer 

107 return pointer 

108 

109 def serialize_frame_set[T: ArchiveTree]( 

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

111 ) -> JsonRef: 

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

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

114 return pointer 

115 

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

117 return iter(self._frame_sets) 

118 

119 def add_array( 

120 self, 

121 array: np.ndarray, 

122 *, 

123 name: str | None = None, 

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

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

126 options_name: str | None = None, 

127 ) -> InlineArrayModel: 

128 return InlineArrayModel( 

129 data=array.tolist(), 

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

131 ) 

132 

133 def add_table( 

134 self, 

135 table: astropy.table.Table, 

136 *, 

137 name: str | None = None, 

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

139 ) -> TableModel: 

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

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

142 

143 def add_structured_array( 

144 self, 

145 array: np.ndarray, 

146 *, 

147 name: str | None = None, 

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

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

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

151 ) -> TableModel: 

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

153 for c in columns: 

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

155 c.unit = unit 

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

157 c.description = description 

158 return TableModel(columns=columns) 

159 

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

161 """Finish serialization. 

162 

163 Parameters 

164 ---------- 

165 tree 

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

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

168 """ 

169 tree.indirect = self._indirect 

170 return tree