Coverage for python/diagnostic_and_simulation_base/nested_dictionary.py: 24%

1# mypy: ignore-errors 

2# TODO: need to fix mypy errors 

3 

4import typing 

5from collections import OrderedDict 

6 

7import numpy as np 

8 

9 

10class ListOfNestedDict: 

11 def __init__(self, nested_dict_data: "NestedDict") -> None: 

12 """This passes in the entire dictionary""" 

13 self.nested_dict_data = nested_dict_data 

14 

15 def __getitem__(self, key_in: str) -> typing.Any: 

16 # Store the keys, until we get to the "data level" 

17 

18 # print(key_in) 

19 if not hasattr(self, "keys"): 

20 # initialising 

21 self.keys = [key_in] 

22 else: 

23 # Store key 

24 self.keys.append(key_in) 

25 # print('starting __getitem__') 

26 # print(self.keys) 

27 # print(' ') 

28 

29 accumulated_results = [self.nested_dict_data[wild_card_key] for wild_card_key in self.nested_dict_data] 

30 # print('accumulated results') 

31 result_final_level = accumulated_results[0] 

32 for key_now in self.keys: 

33 result_final_level = result_final_level[key_now] 

34 # print('got keys') 

35 

36 # print(type(result_final_level)) 

37 if isinstance(result_final_level, NestedDict): 

38 return self 

39 elif isinstance(result_final_level, ListOfNestedDict): 

40 # keep going recursively, storing "self.keys" until we get to the data 

41 return self 

42 else: 

43 result_to_return = [] 

44 for accumulated_result in accumulated_results: 

45 result_final_level = accumulated_result 

46 for key_now in self.keys: 

47 result_final_level = result_final_level[key_now] 

48 result_to_return.append(result_final_level) 

49 

50 result_to_return_np = np.array(result_to_return) 

51 # move the 0th axis to the end - this gives the desired shape 

52 new_axes = tuple(range(1, result_to_return_np.ndim)) + (0,) 

53 result_to_return_np = np.transpose(result_to_return_np, axes=new_axes) 

54 return result_to_return_np 

55 

56 # def __len__(self): 

57 # return len(self.nested_dict_data.keys()) # this probably won't work??? 

58 

59 

60class NestedDict(OrderedDict): 

61 """A custom nested dictionary 

62 

63 Wild card accumulator: In MDSplus, and all of our code, by convention time is the 0th index. 

64 If there is vecor_data inside a node, then it is highly likely that it is a time-dependent quantity 

65 

66 Example, showing the array ordering: 

67 ```python 

68 from diagnostics_analysis_base import NestedDict 

69 import numpy as np 

70 

71 x = NestedDict() 

72 

73 x["top_level_1"]["middle_level_1"]["bottom_level_1"] = np.random.rand(6, 9) 

74 x["top_level_1"]["middle_level_1"]["bottom_level_2"] = np.random.rand(6, 9) 

75 x["top_level_1"]["middle_level_2"]["bottom_level_1"] = np.random.rand(6, 9) 

76 x["top_level_1"]["middle_level_2"]["bottom_level_2"] = np.random.rand(6, 9) 

77 x["top_level_1"]["middle_level_3"]["bottom_level_1"] = np.random.rand(6, 9) 

78 x["top_level_1"]["middle_level_3"]["bottom_level_2"] = np.random.rand(6, 9) 

79 x["top_level_2"]["middle_level_1"]["bottom_level_1"] = np.random.rand(6, 9) 

80 x["top_level_2"]["middle_level_1"]["bottom_level_2"] = np.random.rand(6, 9) 

81 x["top_level_2"]["middle_level_2"]["bottom_level_1"] = np.random.rand(6, 9) 

82 x["top_level_2"]["middle_level_2"]["bottom_level_2"] = np.random.rand(6, 9) 

83 x["top_level_2"]["middle_level_3"]["bottom_level_1"] = np.random.rand(6, 9) 

84 x["top_level_2"]["middle_level_3"]["bottom_level_2"] = np.random.rand(6, 9) 

85 x["top_level_3"]["middle_level_1"]["bottom_level_1"] = np.random.rand(6, 9) 

86 x["top_level_3"]["middle_level_1"]["bottom_level_2"] = np.random.rand(6, 9) 

87 x["top_level_3"]["middle_level_2"]["bottom_level_1"] = np.random.rand(6, 9) 

88 x["top_level_3"]["middle_level_2"]["bottom_level_2"] = np.random.rand(6, 9) 

89 x["top_level_3"]["middle_level_3"]["bottom_level_1"] = np.random.rand(6, 9) 

90 x["top_level_3"]["middle_level_3"]["bottom_level_2"] = np.random.rand(6, 9) 

91 x["top_level_4"]["middle_level_1"]["bottom_level_1"] = np.random.rand(6, 9) 

92 x["top_level_4"]["middle_level_1"]["bottom_level_2"] = np.random.rand(6, 9) 

93 x["top_level_4"]["middle_level_2"]["bottom_level_1"] = np.random.rand(6, 9) 

94 x["top_level_4"]["middle_level_2"]["bottom_level_2"] = np.random.rand(6, 9) 

95 x["top_level_4"]["middle_level_3"]["bottom_level_1"] = np.random.rand(6, 9) 

96 x["top_level_4"]["middle_level_3"]["bottom_level_2"] = np.random.rand(6, 9) 

97 

98 y = x["*"]["*"]["*"] 

99 assert y.shape == (6, 9, 2, 3, 4) 

100 ``` 

101 """ 

102 

103 def __missing__(self, key: str) -> "NestedDict": 

104 """Add to the missing dict""" 

105 value = self[key] = type(self)() 

106 return value 

107 

108 def to_dictionary(self) -> dict[str, typing.Any]: 

109 """Convert NestedDict to a standard python `dict`""" 

110 result = {} 

111 for key, value in self.items(): 

112 if isinstance(value, NestedDict): 

113 result[key] = value.to_dictionary() 

114 else: 

115 result[key] = value 

116 return result 

117 

118 def __setitem__( 

119 self, 

120 key: str, 

121 value: typing.Any, 

122 ) -> None: 

123 """Override __setitem__ to ensure all dictionaries are converted to NestedDict.""" 

124 value = _convert_to_nested_dict(value) 

125 # TODO: figure out why this fails mypy tests? 

126 OrderedDict.__setitem__(self, key, value) 

127 

128 def update( 

129 self, 

130 dict_to_add: dict[typing.Any, typing.Any] | None = None, 

131 ) -> None: 

132 """Recursively updates this dictionary with another dictionary, 

133 ensuring that nested dictionaries are converted to NestedDict.""" 

134 

135 if dict_to_add is None: 

136 dict_to_add = {} 

137 

138 for key, value in dict_to_add.items(): 

139 if isinstance(value, dict): 

140 # If the key exists and is already a NestedDict, update recursively 

141 if key in self and isinstance(self[key], NestedDict): 

142 self[key].update(value) 

143 else: 

144 # Otherwise, convert the value to NestedDict and assign it 

145 self[key] = NestedDict(value) 

146 else: 

147 # Set the value directly if it's not a dictionary 

148 self[key] = value 

149 

150 # def update( 

151 # self, 

152 # dict_to_add: dict | None = None, 

153 # ) -> None: 

154 # """Recursively updates this dictionary with another dictionary or kwargs, 

155 # ensuring that nested dictionaries are converted to NestedDict.""" 

156 

157 # if dict_to_add is None: 

158 # dict_to_add = {} 

159 

160 # nested_dict_to_add = _convert_to_nested_dict(dict_to_add) 

161 

162 # for key, value in dict_to_add.items(): 

163 # if isinstance(value, dict): 

164 # # Recursively update or set as NestedDict 

165 # self[key] = _convert_to_nested_dict(value) 

166 # else: 

167 # self[key] = value 

168 

169 def print_data(self, indent: int = 0) -> str: 

170 """Recursively generates JSON-like string with proper formatting.""" 

171 spacing = " " * indent # 2 spaces per level of indentation 

172 items = [] 

173 for key, value in self.items(): 

174 if isinstance(value, NestedDict): 

175 # Recursively generate nested dicts 

176 items.append(f'{spacing} "{key}": {{') 

177 items.append(value.print_data(indent + 1)) # Recursive call for nested dict 

178 items.append(f"{spacing} }}") 

179 else: 

180 # Handle other types, using repr() for proper formatting 

181 items.append(f'{spacing} "{key}": {repr(value)}') 

182 

183 return "\n".join(items) # Return the joined string without leading/trailing newlines 

184 

185 def __str__(self) -> str: 

186 """Override __str__ to use print_data for JSON-like string representation.""" 

187 return "{\n" + self.print_data() + "\n}" 

188 

189 def __getitem__(self, key: str) -> typing.Any: 

190 # wild card 

191 if key == "*": 

192 # use wild_card and retrieve the data 

193 accumulated_results = [self[wild_card_key] for wild_card_key in self] 

194 # test if data is at this level or not 

195 if isinstance(accumulated_results[0], NestedDict): 

196 return ListOfNestedDict(self) 

197 else: 

198 # Stack the arrays along a new dimension (axis=0) 

199 stacked_results = np.stack(accumulated_results, axis=-1) # Stack along the last axis 

200 return stacked_results 

201 

202 # Normal behaviour 

203 # TODO: this fails `ty` type checking 

204 return OrderedDict.__getitem__(self, key) 

205 

206 # TODO: need to fix mypy tests 

207 def print_keys(self, d: None = None, path: list[str] | None = None) -> None: 

208 if d is None: 

209 # TODO: this fails `ty` type checking. I think this is because we are re-defining `d` from type None to type NestedDict 

210 d = self # Use the current instance as the dictionary 

211 

212 if path is None: 

213 path = [] 

214 

215 # Check if keys_long is already initialized in self 

216 if not hasattr(self, "keys_long"): 

217 self.keys_long = [] 

218 self.data_type_long = [] 

219 

220 # TODO: this fails `ty` type checking 

221 for key, value in d.items(): 

222 new_path = path + [f'["{key}"]'] 

223 

224 if isinstance(value, dict): 

225 # Recursively handle nested dictionaries 

226 # print("BUXTON: error") 

227 self.print_keys(value, new_path) 

228 else: 

229 # Prepare the type information 

230 if isinstance(value, np.ndarray): 

231 type_info = f" = np.ndarray; shape={value.shape}" 

232 elif isinstance(value, list): 

233 type_info = f" = list; length={len(value)}" 

234 else: 

235 type_info = f" = {type(value).__name__}" 

236 

237 # Print the path with the type info 

238 self.keys_long.append("".join(new_path)) # Store key path 

239 self.data_type_long.append(type_info) # Store type info 

240 

241 if path == []: # Check if we are at the top level 

242 # Find the length of the longest key path 

243 max_key_length = max(len(key) for key in self.keys_long) 

244 

245 # Print with proper alignment 

246 for key, type_info in zip(self.keys_long, self.data_type_long): 

247 # Calculate required spaces for alignment 

248 spaces = " " * (max_key_length - len(key) + 4) # 4 spaces for padding 

249 print(f"{key}{spaces}{type_info}") 

250 

251 

252# BUXTON: is "typing.Any" only option?? 

253def _convert_to_nested_dict(value: typing.Any) -> typing.Any: 

254 """Helper function to recursively convert all dictionaries to NestedDict.""" 

255 if isinstance(value, OrderedDict) and not isinstance(value, NestedDict): 

256 # Create a new NestedDict and populate it using a for loop 

257 nested = NestedDict() 

258 for k, v in value.items(): 

259 nested[k] = _convert_to_nested_dict(v) # Recursively convert 

260 return nested 

261 return value