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

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", "read") 

15 

16from collections.abc import Callable 

17from types import EllipsisType 

18from typing import TYPE_CHECKING, Any 

19 

20import astropy.table 

21import numpy as np 

22 

23from lsst.resources import ResourcePath, ResourcePathExpression 

24 

25from .._transforms import FrameSet 

26from ..serialization import ( 

27 ArchiveReadError, 

28 ArchiveTree, 

29 ArrayReferenceModel, 

30 InlineArrayModel, 

31 InputArchive, 

32 JsonRef, 

33 ReadResult, 

34 TableModel, 

35 no_header_updates, 

36) 

37 

38if TYPE_CHECKING: 

39 import astropy.io.fits 

40 

41 

42def read[T: Any]( 

43 cls: type[T], 

44 target: ResourcePathExpression | ArchiveTree, 

45 **kwargs: Any, 

46) -> ReadResult[T]: 

47 """Read an object from a JSON file. 

48 

49 Parameters 

50 ---------- 

51 target 

52 File to read (convertible to `lsst.resources.ResourcePath`) or an 

53 `.serialization.ArchiveTree` to finish deserializing. If the latter, 

54 its ``indirect`` `list` will be interpreted and then cleared. 

55 **kwargs 

56 Extra keyword arguments passed to ``cls.deserialize`` (e.g. ``bbox`` 

57 for image subset reads), matching the FITS and NDF backends. 

58 

59 Returns 

60 ------- 

61 ReadResult 

62 A named tuple containing the deserialized object and any additional 

63 metadata or butler information saved alongside it. 

64 

65 Notes 

66 ----- 

67 Supported types must implement ``deserialize`` and 

68 ``_get_archive_tree_type`` (see `.Image` for an example). 

69 """ 

70 tree_type: type[ArchiveTree] = cls._get_archive_tree_type(JsonRef) 

71 if not isinstance(target, ArchiveTree): 

72 target = tree_type.model_validate_json(ResourcePath(target).read()) 

73 archive = JsonInputArchive(target.indirect) 

74 obj: T = target.deserialize(archive, **kwargs) 

75 target.indirect = [] 

76 return ReadResult(obj, target.metadata, target.butler_info) 

77 

78 

79class JsonInputArchive(InputArchive[JsonRef]): 

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

81 reads from JSON files. 

82 

83 Parameters 

84 ---------- 

85 indirect 

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

87 serialization model. 

88 """ 

89 

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

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

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

93 

94 def deserialize_pointer[U: ArchiveTree, V]( 

95 self, 

96 pointer: JsonRef, 

97 model_type: type[U], 

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

99 ) -> V: 

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

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

102 return existing 

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

104 result = deserializer(model, self) 

105 self._deserialized_pointer_cache[index] = result 

106 return result 

107 

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

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

110 try: 

111 result = self._deserialized_pointer_cache[index] 

112 except KeyError: 

113 raise AssertionError( 

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

115 "before any dependent transform can be." 

116 ) from None 

117 if not isinstance(result, FrameSet): 

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

119 return result 

120 

121 def get_array( 

122 self, 

123 model: ArrayReferenceModel | InlineArrayModel, 

124 *, 

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

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

127 ) -> np.ndarray: 

128 if not isinstance(model, InlineArrayModel): 

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

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

131 

132 def get_table( 

133 self, 

134 model: TableModel, 

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

136 ) -> astropy.table.Table: 

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

138 for column_model in model.columns: 

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

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

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

142 column_model.data.data, 

143 name=column_model.name, 

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

145 unit=column_model.unit, 

146 description=column_model.description, 

147 meta=column_model.meta, 

148 ) 

149 return result 

150 

151 def get_structured_array( 

152 self, 

153 model: TableModel, 

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

155 ) -> np.ndarray: 

156 table = self.get_table(model) 

157 return table.as_array()