Coverage for python/lsst/images/json/_input_archive.py: 68%

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__ = ("JsonInputArchive",) 

15 

16from collections.abc import Callable, Iterator 

17from contextlib import contextmanager 

18from types import EllipsisType 

19from typing import TYPE_CHECKING, Any, Self 

20 

21import astropy.table 

22import numpy as np 

23from pydantic_core import from_json 

24 

25from lsst.resources import ResourcePath, ResourcePathExpression 

26 

27from .._transforms import FrameSet 

28from ..serialization import ( 

29 ArchiveInfo, 

30 ArchiveReadError, 

31 ArchiveTree, 

32 ArrayReferenceModel, 

33 InlineArrayModel, 

34 InputArchive, 

35 JsonRef, 

36 TableModel, 

37 no_header_updates, 

38 parameterize_tree, 

39 tree_class_for_info, 

40) 

41 

42if TYPE_CHECKING: 

43 import astropy.io.fits 

44 

45 

46class JsonInputArchive(InputArchive[JsonRef]): 

47 """An implementation of the `.serialization.InputArchive` interface that 

48 reads from JSON files. 

49 

50 Parameters 

51 ---------- 

52 indirect 

53 The `.serialization.ArchiveTree.indirect` attribute of the root 

54 serialization model. 

55 """ 

56 

57 @classmethod 

58 def get_basic_info(cls, path: ResourcePathExpression) -> ArchiveInfo: 

59 """Read the top-level tree's ``schema_url``; JSON has no container 

60 format version. 

61 

62 This parses the whole document. Unlike the FITS and NDF backends 

63 there is no cheap header to read: ``schema_url`` is a computed field 

64 serialized after the (potentially large) ``indirect`` payload, and 

65 nested trees carry their own ``schema_url``, so a bounded prefix 

66 cannot identify the top-level tree reliably. JSON is not intended 

67 for large pixel archives, where FITS or NDF should be used instead. 

68 

69 Parameters 

70 ---------- 

71 path 

72 Path to the archive to read. 

73 """ 

74 raw = from_json(ResourcePath(path).read()) 

75 if not isinstance(raw, dict) or not raw.get("schema_url"): 75 ↛ 76line 75 didn't jump to line 76 because the condition on line 75 was never true

76 raise ArchiveReadError(f"{path!r} has no schema_url in its top-level JSON tree.") 

77 return ArchiveInfo.from_schema_url(raw["schema_url"], format_version=None) 

78 

79 @classmethod 

80 @contextmanager 

81 def open_tree( 

82 cls, 

83 path: ResourcePathExpression, 

84 *, 

85 partial: bool = True, 

86 **backend_kwargs: Any, 

87 ) -> Iterator[tuple[Self, ArchiveTree, ArchiveInfo]]: 

88 """Parse the JSON tree and yield ``(archive, tree, info)``. 

89 

90 Parameters 

91 ---------- 

92 path 

93 File resource to open. 

94 partial 

95 Ignored. The entire JSON file is always read into memory. 

96 **backend_kwargs 

97 No keyword parameters are supported by this backend. 

98 """ 

99 raw = ResourcePath(path).read() 

100 parsed = from_json(raw) 

101 if not isinstance(parsed, dict) or not parsed.get("schema_url"): 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true

102 raise ArchiveReadError(f"{path!r} has no schema_url in its top-level JSON tree.") 

103 info = ArchiveInfo.from_schema_url(parsed["schema_url"], format_version=None) 

104 tree_cls = tree_class_for_info(info, path) 

105 parameterized = parameterize_tree(tree_cls, JsonRef) 

106 tree = parameterized.model_validate_json(raw) 

107 archive = cls(tree.indirect) 

108 try: 

109 yield archive, tree, info 

110 finally: 

111 tree.indirect = [] 

112 

113 def __init__(self, indirect: list[Any] | None = None) -> None: 

114 self._indirect = indirect if indirect is not None else [] 

115 self._deserialized_pointer_cache: dict[int, Any] = {} 

116 

117 def deserialize_pointer[U: ArchiveTree, V]( 

118 self, 

119 pointer: JsonRef, 

120 model_type: type[U], 

121 deserializer: Callable[[U, InputArchive[JsonRef]], V], 

122 ) -> V: 

123 index = int(pointer.ref.removeprefix("#/indirect/")) 

124 if (existing := self._deserialized_pointer_cache.get(index)) is not None: 

125 return existing 

126 model = model_type.model_validate(self._indirect[index]) 

127 result = deserializer(model, self) 

128 self._deserialized_pointer_cache[index] = result 

129 return result 

130 

131 def get_frame_set(self, ref: JsonRef) -> FrameSet: 

132 index = int(ref.ref.removeprefix("#/indirect/")) 

133 try: 

134 result = self._deserialized_pointer_cache[index] 

135 except KeyError: 

136 raise AssertionError( 

137 f"Frame set at {ref.model_dump_json(indent=2)} must be deserialized " 

138 "before any dependent transform can be." 

139 ) from None 

140 if not isinstance(result, FrameSet): 

141 raise ArchiveReadError(f"Expected a FrameSet instance at {ref.model_dump_json(indent=2)}.") 

142 return result 

143 

144 def get_array( 

145 self, 

146 model: ArrayReferenceModel | InlineArrayModel, 

147 *, 

148 slices: tuple[slice, ...] | EllipsisType = ..., 

149 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

150 ) -> np.ndarray: 

151 if not isinstance(model, InlineArrayModel): 151 ↛ 152line 151 didn't jump to line 152 because the condition on line 151 was never true

152 raise ArchiveReadError("Only inline arrays are supported in JSON archives.") 

153 return np.array(model.data, dtype=model.datatype.to_numpy())[slices] 

154 

155 def get_table( 

156 self, 

157 model: TableModel, 

158 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

159 ) -> astropy.table.Table: 

160 result = astropy.table.Table(meta=model.meta) 

161 for column_model in model.columns: 

162 if not isinstance(column_model.data, InlineArrayModel): 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true

163 raise ArchiveReadError("Only inline arrays are supported in JSON archives.") 

164 result[column_model.name] = astropy.table.Column( 

165 column_model.data.data, 

166 name=column_model.name, 

167 dtype=column_model.data.datatype.to_numpy(), 

168 unit=column_model.unit, 

169 description=column_model.description, 

170 meta=column_model.meta, 

171 ) 

172 return result 

173 

174 def get_structured_array( 

175 self, 

176 model: TableModel, 

177 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

178 ) -> np.ndarray: 

179 table = self.get_table(model) 

180 return table.as_array()