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](cls: type[T], target: ResourcePathExpression | ArchiveTree) -> ReadResult[T]: 

43 """Read an object from a FITS file. 

44 

45 Parameters 

46 ---------- 

47 target 

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

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

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

51 

52 Returns 

53 ------- 

54 ReadResult 

55 A named tuple containing the deserialized object and any additional 

56 metadata or butler information saved alongside it. 

57 

58 Notes 

59 ----- 

60 Supported types must implement ``deserialize`` and 

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

62 """ 

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

64 if not isinstance(target, ArchiveTree): 

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

66 archive = JsonInputArchive(target.indirect) 

67 obj = cls.deserialize(target, archive) 

68 target.indirect = [] 

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

70 

71 

72class JsonInputArchive(InputArchive[JsonRef]): 

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

74 reads from JSON files. 

75 

76 Parameters 

77 ---------- 

78 indirect 

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

80 serialization model. 

81 """ 

82 

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

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

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

86 

87 def deserialize_pointer[U: ArchiveTree, V]( 

88 self, 

89 pointer: JsonRef, 

90 model_type: type[U], 

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

92 ) -> V: 

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

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

95 return existing 

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

97 result = deserializer(model, self) 

98 self._deserialized_pointer_cache[index] = result 

99 return result 

100 

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

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

103 try: 

104 result = self._deserialized_pointer_cache[index] 

105 except KeyError: 

106 raise AssertionError( 

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

108 "before any dependent transform can be." 

109 ) from None 

110 if not isinstance(result, FrameSet): 

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

112 return result 

113 

114 def get_array( 

115 self, 

116 model: ArrayReferenceModel | InlineArrayModel, 

117 *, 

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

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

120 ) -> np.ndarray: 

121 if not isinstance(model, InlineArrayModel): 

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

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

124 

125 def get_table( 

126 self, 

127 model: TableModel, 

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

129 ) -> astropy.table.Table: 

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

131 for column_model in model.columns: 

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

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

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

135 column_model.data.data, 

136 name=column_model.name, 

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

138 unit=column_model.unit, 

139 description=column_model.description, 

140 meta=column_model.meta, 

141 ) 

142 return result 

143 

144 def get_structured_array( 

145 self, 

146 model: TableModel, 

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

148 ) -> np.ndarray: 

149 table = self.get_table(model) 

150 return table.as_array()