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

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) 

40 

41if TYPE_CHECKING: 

42 import astropy.io.fits 

43 

44 

45class JsonInputArchive(InputArchive[JsonRef]): 

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

47 reads from JSON files. 

48 

49 Parameters 

50 ---------- 

51 indirect 

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

53 serialization model. 

54 """ 

55 

56 @classmethod 

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

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

59 format version. 

60 

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

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

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

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

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

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

67 """ 

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

69 if not isinstance(raw, dict) or not raw.get("schema_url"): 

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

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

72 

73 @classmethod 

74 @contextmanager 

75 def open_tree( 

76 cls, 

77 path: ResourcePathExpression, 

78 tree_cls: type[ArchiveTree], 

79 *, 

80 partial: bool = True, 

81 **backend_kwargs: Any, 

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

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

84 

85 A no-resource context manager: JSON is fully in memory, so 

86 ``partial`` is a no-op. 

87 ``tree.indirect`` is released when the context exits. 

88 """ 

89 parameterized = parameterize_tree(tree_cls, JsonRef) 

90 tree = parameterized.model_validate_json(ResourcePath(path).read()) 

91 archive = cls(tree.indirect) 

92 try: 

93 yield archive, tree 

94 finally: 

95 tree.indirect = [] 

96 

97 def __init__(self, indirect: list[Any] | None = None): 

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

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

100 

101 def deserialize_pointer[U: ArchiveTree, V]( 

102 self, 

103 pointer: JsonRef, 

104 model_type: type[U], 

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

106 ) -> V: 

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

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

109 return existing 

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

111 result = deserializer(model, self) 

112 self._deserialized_pointer_cache[index] = result 

113 return result 

114 

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

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

117 try: 

118 result = self._deserialized_pointer_cache[index] 

119 except KeyError: 

120 raise AssertionError( 

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

122 "before any dependent transform can be." 

123 ) from None 

124 if not isinstance(result, FrameSet): 

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

126 return result 

127 

128 def get_array( 

129 self, 

130 model: ArrayReferenceModel | InlineArrayModel, 

131 *, 

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

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

134 ) -> np.ndarray: 

135 if not isinstance(model, InlineArrayModel): 

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

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

138 

139 def get_table( 

140 self, 

141 model: TableModel, 

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

143 ) -> astropy.table.Table: 

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

145 for column_model in model.columns: 

146 if not isinstance(column_model.data, InlineArrayModel): 

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

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

149 column_model.data.data, 

150 name=column_model.name, 

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

152 unit=column_model.unit, 

153 description=column_model.description, 

154 meta=column_model.meta, 

155 ) 

156 return result 

157 

158 def get_structured_array( 

159 self, 

160 model: TableModel, 

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

162 ) -> np.ndarray: 

163 table = self.get_table(model) 

164 return table.as_array()