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

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): 

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

35 elif np.issubdtype(arr.dtype, np.floating): 

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: 

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: 

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: 

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

153 if key not in metadata: 

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

155 return metadata[key]