Coverage for python/lsst/scarlet/lite/io/utils.py: 78%

1from typing import Any 

2 

3import numpy as np 

4 

5__all__ = ["PersistenceError"] 

6 

7 

8class PersistenceError(Exception): 

9 """Custom error for persistence issues.""" 

10 

11 pass 

12 

13 

14def numpy_to_json(arr: np.ndarray) -> dict[str, Any]: 

15 """ 

16 Encode a numpy array as JSON-serializable dictionary. 

17 

18 Parameters 

19 ---------- 

20 arr : 

21 The numpy array to encode 

22 

23 Returns 

24 ------- 

25 result : 

26 A JSON formatted dictionary containing the dtype, shape, 

27 and data of the array. 

28 """ 

29 # Convert to native Python types for JSON serialization 

30 flattened = arr.flatten() 

31 

32 # Convert numpy scalars to native Python types 

33 if np.issubdtype(arr.dtype, np.integer): 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 data: list = [int(x) for x in flattened] 

35 elif np.issubdtype(arr.dtype, np.floating): 35 ↛ 37line 35 didn't jump to line 37 because the condition on line 35 was always true

36 data = [float(x) for x in flattened] 

37 elif np.issubdtype(arr.dtype, np.complexfloating): 

38 data = [complex(x) for x in flattened] 

39 elif np.issubdtype(arr.dtype, np.bool_): 

40 data = [bool(x) for x in flattened] 

41 else: 

42 # For other types (strings, objects, etc.), convert to string 

43 data = [str(x) for x in flattened] 

44 

45 return {"dtype": str(arr.dtype), "shape": tuple(arr.shape), "data": data} 

46 

47 

48def json_to_numpy(encoded_dict: dict[str, Any]) -> np.ndarray: 

49 """ 

50 Decode a JSON dictionary back to a numpy array. 

51 

52 Parameters 

53 ---------- 

54 encoded_dict : 

55 Dictionary with 'dtype', 'shape', and 'data' keys. 

56 

57 Returns 

58 ------- 

59 result : 

60 The reconstructed numpy array. 

61 """ 

62 if "dtype" not in encoded_dict or "shape" not in encoded_dict or "data" not in encoded_dict: 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true

63 raise ValueError("Encoded dictionary must contain 'dtype', 'shape', and 'data' keys.") 

64 return np.array(encoded_dict["data"], dtype=encoded_dict["dtype"]).reshape(encoded_dict["shape"]) 

65 

66 

67def encode_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] | None: 

68 """Pack metadata into a JSON compatible format. 

69 

70 Parameters 

71 ---------- 

72 metadata : 

73 The metadata to be packed. 

74 

75 Returns 

76 ------- 

77 result : 

78 The packed metadata. 

79 """ 

80 if metadata is None: 

81 return None 

82 encoded = {} 

83 array_keys = [] 

84 for key, value in metadata.items(): 

85 if isinstance(value, np.ndarray): 

86 _encoded = numpy_to_json(value) 

87 encoded[key] = _encoded["data"] 

88 encoded[f"{key}_shape"] = _encoded["shape"] 

89 encoded[f"{key}_dtype"] = _encoded["dtype"] 

90 array_keys.append(key) 

91 else: 

92 encoded[key] = value 

93 if len(array_keys) > 0: 

94 encoded["array_keys"] = array_keys 

95 return encoded 

96 

97 

98def decode_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] | None: 

99 """Unpack metadata from a JSON compatible format. 

100 

101 Parameters 

102 ---------- 

103 metadata : 

104 The metadata to be unpacked. 

105 

106 Returns 

107 ------- 

108 result : 

109 The unpacked metadata. 

110 """ 

111 if metadata is None: 

112 return None 

113 if "array_keys" in metadata: 

114 for key in metadata["array_keys"]: 

115 # Default dtype is float32 to support legacy models 

116 dtype = metadata.pop(f"{key}_dtype", "float32") 

117 shape = metadata.pop(f"{key}_shape", None) 

118 if shape is None and f"{key}Shape" in metadata: 118 ↛ 120line 118 didn't jump to line 120 because the condition on line 118 was never true

119 # Support legacy models that use `keyShape` 

120 shape = metadata[f"{key}Shape"] 

121 decoded = json_to_numpy({"dtype": dtype, "shape": shape, "data": metadata[key]}) 

122 metadata[key] = decoded 

123 # Remove the array keys after decoding 

124 del metadata["array_keys"] 

125 return metadata 

126 

127 

128def extract_from_metadata( 

129 data: Any, 

130 metadata: dict[str, Any] | None, 

131 key: str, 

132) -> Any: 

133 """Extract relevant information from the metadata. 

134 

135 Parameters 

136 ---------- 

137 data : 

138 The data to extract information from. 

139 metadata : 

140 The metadata to extract information from. 

141 key : 

142 The key to extract from the metadata. 

143 

144 Returns 

145 ------- 

146 result : 

147 A tuple containing the extracted data and metadata. 

148 """ 

149 if data is not None: 

150 return data 

151 if metadata is None: 151 ↛ 152line 151 didn't jump to line 152 because the condition on line 151 was never true

152 raise ValueError("Both data and metadata cannot be None") 

153 if key not in metadata: 153 ↛ 154line 153 didn't jump to line 154 because the condition on line 153 was never true

154 raise ValueError(f"'{key}' not found in metadata") 

155 return metadata[key]