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 raw = from_json(ResourcePath(path).read()) 

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

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

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

73 

74 @classmethod 

75 @contextmanager 

76 def open_tree( 

77 cls, 

78 path: ResourcePathExpression, 

79 *, 

80 partial: bool = True, 

81 **backend_kwargs: Any, 

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

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

84 

85 Parameters 

86 ---------- 

87 path 

88 File resource to open. 

89 partial 

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

91 **backend_kwargs 

92 No keyword parameters are supported by this backend. 

93 """ 

94 raw = ResourcePath(path).read() 

95 parsed = from_json(raw) 

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

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

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

99 tree_cls = tree_class_for_info(info, path) 

100 parameterized = parameterize_tree(tree_cls, JsonRef) 

101 tree = parameterized.model_validate_json(raw) 

102 archive = cls(tree.indirect) 

103 try: 

104 yield archive, tree, info 

105 finally: 

106 tree.indirect = [] 

107 

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

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

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

111 

112 def deserialize_pointer[U: ArchiveTree, V]( 

113 self, 

114 pointer: JsonRef, 

115 model_type: type[U], 

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

117 ) -> V: 

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

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

120 return existing 

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

122 result = deserializer(model, self) 

123 self._deserialized_pointer_cache[index] = result 

124 return result 

125 

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

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

128 try: 

129 result = self._deserialized_pointer_cache[index] 

130 except KeyError: 

131 raise AssertionError( 

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

133 "before any dependent transform can be." 

134 ) from None 

135 if not isinstance(result, FrameSet): 

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

137 return result 

138 

139 def get_array( 

140 self, 

141 model: ArrayReferenceModel | InlineArrayModel, 

142 *, 

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

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

145 ) -> np.ndarray: 

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

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

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

149 

150 def get_table( 

151 self, 

152 model: TableModel, 

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

154 ) -> astropy.table.Table: 

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

156 for column_model in model.columns: 

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

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

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

160 column_model.data.data, 

161 name=column_model.name, 

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

163 unit=column_model.unit, 

164 description=column_model.description, 

165 meta=column_model.meta, 

166 ) 

167 return result 

168 

169 def get_structured_array( 

170 self, 

171 model: TableModel, 

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

173 ) -> np.ndarray: 

174 table = self.get_table(model) 

175 return table.as_array()