Coverage for muutils / json_serialize / util.py: 50%

1"""utilities for json_serialize""" 

2 

3from __future__ import annotations 

4 

5import dataclasses 

6import functools 

7import inspect 

8import sys 

9import typing 

10import warnings 

11from typing import Any, Callable, Iterable, TypeVar, Union 

12 

13from muutils.json_serialize.types import BaseType, Hashableitem 

14 

15if typing.TYPE_CHECKING: 

16 pass 

17 

18_NUMPY_WORKING: bool 

19try: 

20 _NUMPY_WORKING = True 

21except ImportError: 

22 warnings.warn("numpy not found, cannot serialize numpy arrays!") 

23 _NUMPY_WORKING = False 

24 

25 

26# pyright: reportExplicitAny=false 

27 

28# At type-checking time, include array serialization types to avoid nominal type errors 

29# This avoids superfluous imports at runtime 

30# if TYPE_CHECKING: 

31# from muutils.json_serialize.array import NumericList, SerializedArrayWithMeta 

32 

33# JSONitem = Union[ 

34# BaseType, 

35# typing.Sequence["JSONitem"], 

36# typing.Dict[str, "JSONitem"], 

37# SerializedArrayWithMeta, 

38# NumericList, 

39# ] 

40# else: 

41 

42JSONitem = Union[ 

43 BaseType, 

44 typing.Sequence["JSONitem"], 

45 typing.Dict[str, "JSONitem"], 

46 # TODO: figure this out 

47 # "_SerializedSet", 

48 # "_SerializedFrozenset", 

49] 

50 

51JSONdict = typing.Dict[str, JSONitem] 

52 

53 

54# TODO: this bit is very broken 

55# or if python version <3.9 

56if typing.TYPE_CHECKING or sys.version_info < (3, 9): 

57 MonoTuple = typing.Sequence 

58else: 

59 

60 class MonoTuple: # pyright: ignore[reportUnreachable] 

61 """tuple type hint, but for a tuple of any length with all the same type""" 

62 

63 __slots__ = () 

64 

65 def __new__(cls, *args, **kwargs): 

66 raise TypeError("Type MonoTuple cannot be instantiated.") 

67 

68 def __init_subclass__(cls, *args, **kwargs): 

69 raise TypeError(f"Cannot subclass {cls.__module__}") 

70 

71 # idk why mypy thinks there is no such function in typing 

72 @typing._tp_cache # type: ignore 

73 def __class_getitem__(cls, params): 

74 if getattr(params, "__origin__", None) == typing.Union: 

75 return typing.GenericAlias(tuple, (params, Ellipsis)) 

76 elif isinstance(params, type): 

77 typing.GenericAlias(tuple, (params, Ellipsis)) 

78 # test if has len and is iterable 

79 elif isinstance(params, Iterable): 

80 if len(params) == 0: 

81 return tuple 

82 elif len(params) == 1: 

83 return typing.GenericAlias(tuple, (params[0], Ellipsis)) 

84 else: 

85 raise TypeError(f"MonoTuple expects 1 type argument, got {params = }") 

86 

87 

88# TYPING: we allow `Any` here because the container is... universal 

89class UniversalContainer: 

90 """contains everything -- `x in UniversalContainer()` is always True""" 

91 

92 def __contains__(self, x: Any) -> bool: # pyright: ignore[reportAny] 

93 return True 

94 

95 

96def isinstance_namedtuple(x: Any) -> bool: # pyright: ignore[reportAny] 

97 """checks if `x` is a `namedtuple` 

98 

99 credit to https://stackoverflow.com/questions/2166818/how-to-check-if-an-object-is-an-instance-of-a-namedtuple 

100 """ 

101 t: type = type(x) # pyright: ignore[reportUnknownVariableType, reportAny] 

102 b: tuple[type, ...] = t.__bases__ 

103 if len(b) != 1 or (b[0] is not tuple): 

104 return False 

105 f: Any = getattr(t, "_fields", None) 

106 if not isinstance(f, tuple): 

107 return False 

108 # fine that the type is unknown -- that's what we want to check 

109 return all(isinstance(n, str) for n in f) # pyright: ignore[reportUnknownVariableType] 

110 

111 

112T_FuncTryCatchReturn = TypeVar("T_FuncTryCatchReturn") 

113 

114 

115def try_catch( 

116 func: Callable[..., T_FuncTryCatchReturn], 

117) -> Callable[..., Union[T_FuncTryCatchReturn, str]]: 

118 """wraps the function to catch exceptions, returns serialized error message on exception 

119 

120 returned func will return normal result on success, or error message on exception 

121 """ 

122 

123 @functools.wraps(func) 

124 def newfunc(*args: Any, **kwargs: Any) -> Union[T_FuncTryCatchReturn, str]: # pyright: ignore[reportAny] 

125 try: 

126 return func(*args, **kwargs) 

127 except Exception as e: 

128 return f"{e.__class__.__name__}: {e}" 

129 

130 return newfunc 

131 

132 

133# TYPING: can we get rid of any of these? 

134def _recursive_hashify(obj: Any, force: bool = True) -> Hashableitem: # pyright: ignore[reportAny] 

135 if isinstance(obj, typing.Mapping): 

136 return tuple((k, _recursive_hashify(v)) for k, v in obj.items()) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] 

137 elif isinstance(obj, (bool, int, float, str)): 

138 return obj 

139 elif isinstance(obj, (tuple, list, Iterable)): 

140 return tuple(_recursive_hashify(v) for v in obj) # pyright: ignore[reportUnknownVariableType] 

141 else: 

142 if force: 

143 return str(obj) # pyright: ignore[reportAny] 

144 else: 

145 raise ValueError(f"cannot hashify:\n{obj}") 

146 

147 

148class SerializationException(Exception): 

149 pass 

150 

151 

152def string_as_lines(s: str | None) -> list[str]: 

153 """for easier reading of long strings in json, split up by newlines 

154 

155 sort of like how jupyter notebooks do it 

156 """ 

157 if s is None: 

158 return list() 

159 else: 

160 return s.splitlines(keepends=False) 

161 

162 

163def safe_getsource(func: Callable[..., Any]) -> list[str]: 

164 try: 

165 return string_as_lines(inspect.getsource(func)) 

166 except Exception as e: 

167 return string_as_lines(f"Error: Unable to retrieve source code:\n{e}") 

168 

169 

170# credit to https://stackoverflow.com/questions/51743827/how-to-compare-equality-of-dataclasses-holding-numpy-ndarray-boola-b-raises 

171def array_safe_eq(a: Any, b: Any) -> bool: # pyright: ignore[reportAny] 

172 """check if two objects are equal, account for if numpy arrays or torch tensors""" 

173 if a is b: 

174 return True 

175 

176 if type(a) is not type(b): # pyright: ignore[reportAny] 

177 return False 

178 

179 if ( 

180 str(type(a)) == "<class 'numpy.ndarray'>" # pyright: ignore[reportAny, reportUnknownArgumentType] 

181 and str(type(b)) == "<class 'numpy.ndarray'>" # pyright: ignore[reportAny, reportUnknownArgumentType] 

182 ) or ( 

183 str(type(a)) == "<class 'torch.Tensor'>" # pyright: ignore[reportAny, reportUnknownArgumentType] 

184 and str(type(b)) == "<class 'torch.Tensor'>" # pyright: ignore[reportAny, reportUnknownArgumentType] 

185 ): 

186 return (a == b).all() # pyright: ignore[reportAny] 

187 

188 if ( 

189 str(type(a)) == "<class 'pandas.core.frame.DataFrame'>" # pyright: ignore[reportUnknownArgumentType, reportAny] 

190 and str(type(b)) == "<class 'pandas.core.frame.DataFrame'>" # pyright: ignore[reportUnknownArgumentType, reportAny] 

191 ): 

192 return a.equals(b) # pyright: ignore[reportAny] 

193 

194 if isinstance(a, typing.Sequence) and isinstance(b, typing.Sequence): 

195 if len(a) == 0 and len(b) == 0: 

196 return True 

197 return len(a) == len(b) and all(array_safe_eq(a1, b1) for a1, b1 in zip(a, b)) 

198 

199 if isinstance(a, (dict, typing.Mapping)) and isinstance(b, (dict, typing.Mapping)): 

200 return len(a) == len(b) and all( # pyright: ignore[reportUnknownArgumentType] 

201 array_safe_eq(k1, k2) and array_safe_eq(a[k1], b[k2]) 

202 for k1, k2 in zip(a.keys(), b.keys()) # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] 

203 ) 

204 

205 try: 

206 return bool(a == b) # pyright: ignore[reportAny] 

207 except (TypeError, ValueError) as e: 

208 warnings.warn(f"Cannot compare {a} and {b} for equality\n{e}") 

209 return NotImplemented # type: ignore[return-value] 

210 

211 

212# TYPING: see what can be done about so many `Any`s here 

213def dc_eq( 

214 dc1: Any, # pyright: ignore[reportAny] 

215 dc2: Any, # pyright: ignore[reportAny] 

216 except_when_class_mismatch: bool = False, 

217 false_when_class_mismatch: bool = True, 

218 except_when_field_mismatch: bool = False, 

219) -> bool: 

220 """ 

221 checks if two dataclasses which (might) hold numpy arrays are equal 

222 

223 # Parameters: 

224 

225 - `dc1`: the first dataclass 

226 - `dc2`: the second dataclass 

227 - `except_when_class_mismatch: bool` 

228 if `True`, will throw `TypeError` if the classes are different. 

229 if not, will return false by default or attempt to compare the fields if `false_when_class_mismatch` is `False` 

230 (default: `False`) 

231 - `false_when_class_mismatch: bool` 

232 only relevant if `except_when_class_mismatch` is `False`. 

233 if `True`, will return `False` if the classes are different. 

234 if `False`, will attempt to compare the fields. 

235 - `except_when_field_mismatch: bool` 

236 only relevant if `except_when_class_mismatch` is `False` and `false_when_class_mismatch` is `False`. 

237 if `True`, will throw `AttributeError` if the fields are different. 

238 (default: `False`) 

239 

240 # Returns: 

241 - `bool`: True if the dataclasses are equal, False otherwise 

242 

243 # Raises: 

244 - `TypeError`: if the dataclasses are of different classes 

245 - `AttributeError`: if the dataclasses have different fields 

246 

247 ``` 

248 [START] 

249 ▼ 

250 ┌─────────────┐ 

251 │ dc1 is dc2? │───Yes───► (True) 

252 └──────┬──────┘ 

253 │No 

254 ▼ 

255 ┌───────────────┐ 

256 │ classes match?│───Yes───► [compare field values] ───► (True/False) 

257 └──────┬────────┘ 

258 │No 

259 ▼ 

260 ┌────────────────────────────┐ 

261 │ except_when_class_mismatch?│───Yes───► { raise TypeError } 

262 └─────────────┬──────────────┘ 

263 │No 

264 ▼ 

265 ┌────────────────────────────┐ 

266 │ false_when_class_mismatch? │───Yes───► (False) 

267 └─────────────┬──────────────┘ 

268 │No 

269 ▼ 

270 ┌────────────────────────────┐ 

271 │ except_when_field_mismatch?│───No────► [compare field values] 

272 └─────────────┬──────────────┘ 

273 │Yes 

274 ▼ 

275 ┌───────────────┐ 

276 │ fields match? │───Yes───► [compare field values] 

277 └──────┬────────┘ 

278 │No 

279 ▼ 

280 { raise AttributeError } 

281 ``` 

282 

283 """ 

284 if dc1 is dc2: 

285 return True 

286 

287 if dc1.__class__ is not dc2.__class__: # pyright: ignore[reportAny] 

288 if except_when_class_mismatch: 

289 # if the classes don't match, raise an error 

290 raise TypeError( 

291 f"Cannot compare dataclasses of different classes: `{dc1.__class__}` and `{dc2.__class__}`" # pyright: ignore[reportAny] 

292 ) 

293 if false_when_class_mismatch: 

294 # return False immediately without attempting field comparison 

295 return False 

296 # classes don't match but we'll try to compare fields anyway 

297 if except_when_field_mismatch: 

298 dc1_fields: set[str] = set([fld.name for fld in dataclasses.fields(dc1)]) # pyright: ignore[reportAny] 

299 dc2_fields: set[str] = set([fld.name for fld in dataclasses.fields(dc2)]) # pyright: ignore[reportAny] 

300 fields_match: bool = set(dc1_fields) == set(dc2_fields) 

301 if not fields_match: 

302 # if the fields don't match, raise an error 

303 raise AttributeError( 

304 f"dataclasses {dc1} and {dc2} have different fields: `{dc1_fields}` and `{dc2_fields}`" 

305 ) 

306 

307 return all( 

308 array_safe_eq(getattr(dc1, fld.name), getattr(dc2, fld.name)) # pyright: ignore[reportAny] 

309 for fld in dataclasses.fields(dc1) # pyright: ignore[reportAny] 

310 if fld.compare 

311 )