Coverage for zanj/loading.py: 79%

1 from __future__ import annotations

2

3 import json

4 import threading

5 import typing

6 import zipfile

7 from dataclasses import dataclass

8 from pathlib import Path

9 from typing import Any, Callable

10

11 import numpy as np

12

13 try:

14 import pandas as pd # type: ignore[import]

15

16 pandas_DataFrame = pd.DataFrame # type: ignore[no-redef]

17 except ImportError:

18

19 class pandas_DataFrame: # type: ignore[no-redef]

20 def __init__(self, *args, **kwargs):

21 raise ImportError("cannot load pandas DataFrame, pandas is not installed")

22

23

24 from muutils.errormode import ErrorMode

25 from muutils.json_serialize.array import load_array

26 from muutils.json_serialize.json_serialize import ObjectPath

27 from muutils.json_serialize.util import (

28 _FORMAT_KEY,

29 _REF_KEY,

30 JSONdict,

31 JSONitem,

32 safe_getsource,

33 string_as_lines,

34)

35

36 from zanj.externals import (

37 GET_EXTERNAL_LOAD_FUNC,

38 ZANJ_MAIN,

39 ZANJ_META,

40 ExternalItem,

41 _ZANJ_pre,

42)

43

44# pylint: disable=protected-access, dangerous-default-value

45

46

47 def _populate_externals_error_checking(key, item) -> bool:

48 """checks that the key is valid for the item. returns "True" we need to augment the path by accessing the "data" element"""

49

50 # special case for not fully loaded external item which we still need to populate

51 if isinstance(item, typing.Mapping):

52 if (_FORMAT_KEY in item) and item[_FORMAT_KEY].endswith(":external"):

53 if "data" in item:

54 return True

55 else:

56 raise KeyError(

57 f"expected an external item, but could not find data: {list(item.keys())}",

58 f"{item[_FORMAT_KEY]}, {len(item) = }, {item.get('data', '<EMPTY>') = }",

59 )

60

61 # if it's a list, make sure the key is an int and that it's in range

62 if isinstance(item, typing.Sequence):

63 if not isinstance(key, int):

64 raise TypeError(f"improper type: '{type(key) = }', expected int")

65 if key >= len(item):

66 raise IndexError(f"index out of range: '{key = }', expected < {len(item)}")

67

68 # if it's a dict, make sure that the key is a str and that it's in the dict

69 elif isinstance(item, typing.Mapping):

70 if not isinstance(key, str):

71 raise TypeError(f"improper type: '{type(key) = }', expected str")

72 if key not in item:

73 raise KeyError(f"key not in dict: '{key = }', expected in {item.keys()}")

74

75 # otherwise, raise an error

76 else:

77 raise TypeError(f"improper type: '{type(item) = }', expected dict or list")

78

79 return False

80

81

82@dataclass

83 class LoaderHandler:

84 """handler for loading an object from a json file or a ZANJ archive"""

85

86 # TODO: add a separate "asserts" function?

87 # right now, any asserts must happen in `check` or `load` which is annoying with lambdas

88

89 # (json_data, path) -> whether to use this handler

90 check: Callable[[JSONitem, ObjectPath, _ZANJ_pre], bool]

91 # function to load the object (json_data, path) -> loaded_obj

92 load: Callable[[JSONitem, ObjectPath, _ZANJ_pre], Any]

93 # unique identifier for the handler, saved in __muutils_format__ field

94 uid: str

95 # source package of the handler -- note that this might be overridden by ZANJ

96 source_pckg: str

97 # priority of the handler, defaults are all 0

98 priority: int = 0

99 # description of the handler

100 desc: str = "(no description)"

101

102 def serialize(self) -> JSONdict:

103 """serialize the handler info"""

104 return {

105 # get the code and doc of the check function

106 "check": {

107 "code": safe_getsource(self.check),

108 "doc": string_as_lines(self.check.__doc__),

109 },

110 # get the code and doc of the load function

111 "load": {

112 "code": safe_getsource(self.load),

113 "doc": string_as_lines(self.load.__doc__),

114 },

115 # get the uid, source_pckg, priority, and desc

116 "uid": str(self.uid),

117 "source_pckg": str(self.source_pckg),

118 "priority": int(self.priority),

119 "desc": str(self.desc),

120 }

121

122 @classmethod

123 def from_formattedclass(cls, fc: type, priority: int = 0):

124 """create a loader from a class with `serialize`, `load` methods and `__muutils_format__` attribute"""

125 assert hasattr(fc, "serialize")

126 assert callable(fc.serialize) # type: ignore

127 assert hasattr(fc, "load")

128 assert callable(fc.load) # type: ignore

129 assert hasattr(fc, _FORMAT_KEY)

130 assert isinstance(fc.__muutils_format__, str) # type: ignore

131

132 return cls(

133 check=lambda json_item, path=None, z=None: ( # type: ignore[misc]

134 json_item[_FORMAT_KEY] == fc.__muutils_format__ # type: ignore[attr-defined]

135 ),

136 load=lambda json_item, path=None, z=None: fc.load(json_item, path, z), # type: ignore[misc]

137 uid=fc.__muutils_format__, # type: ignore[attr-defined]

138 source_pckg=str(fc.__module__),

139 priority=priority,

140 desc=f"formatted class loader for {fc.__name__}",

141 )

142

143

144# TODO: how can we type hint this without actually importing torch?

145 def _torch_loaderhandler_load(

146 json_item: JSONitem,

147 path: ObjectPath,

148 z: _ZANJ_pre | None = None,

149):

150 """load a torch tensor from a json item"""

151 try:

152 import torch # type: ignore[import-not-found]

153 from muutils.tensor_utils import TORCH_DTYPE_MAP

154 except ImportError as e:

155 err_msg: str = f"could not import torch, which we need to load the object at {path = }: {json_item = }"

156 raise ImportError(err_msg) from e

157

158 return torch.tensor(

159 load_array(json_item),

160 dtype=TORCH_DTYPE_MAP[json_item["dtype"]], # type: ignore[index, call-overload]

161 )

162

163

164# NOTE: there are type ignores on the loaders, since the type checking should be the responsibility of the check function

165

166 LOADER_MAP_LOCK = threading.Lock()

167

168 LOADER_MAP: dict[str, LoaderHandler] = {

169 lh.uid: lh

170 for lh in [

171 # array external

172 LoaderHandler(

173 check=lambda json_item, path=None, z=None: ( # type: ignore[misc]

174 isinstance(json_item, typing.Mapping)

175 and _FORMAT_KEY in json_item

176 and json_item[_FORMAT_KEY].startswith("numpy.ndarray")

177 # and json_item["data"].dtype.name == json_item["dtype"]

178 # and tuple(json_item["data"].shape) == tuple(json_item["shape"])

179 ),

180 load=lambda json_item, path=None, z=None: np.array( # type: ignore[misc]

181 load_array(json_item), dtype=np.dtype(json_item["dtype"])

182 ),

183 uid="numpy.ndarray",

184 source_pckg="zanj",

185 desc="numpy.ndarray loader",

186 ),

187 LoaderHandler(

188 check=lambda json_item, path=None, z=None: ( # type: ignore[misc]

189 isinstance(json_item, typing.Mapping)

190 and _FORMAT_KEY in json_item

191 and json_item[_FORMAT_KEY].startswith("torch.Tensor")

192 # and json_item["data"].dtype.name == json_item["dtype"]

193 # and tuple(json_item["data"].shape) == tuple(json_item["shape"])

194 ),

195 load=_torch_loaderhandler_load,

196 uid="torch.Tensor",

197 source_pckg="zanj",

198 desc="torch.Tensor loader",

199 ),

200 # pandas

201 LoaderHandler(

202 check=lambda json_item, path=None, z=None: ( # type: ignore[misc]

203 isinstance(json_item, typing.Mapping)

204 and _FORMAT_KEY in json_item

205 and json_item[_FORMAT_KEY].startswith("pandas.DataFrame")

206 and "data" in json_item

207 and isinstance(json_item["data"], typing.Sequence)

208 ),

209 load=lambda json_item, path=None, z=None: pandas_DataFrame( # type: ignore[misc]

210 json_item["data"]

211 ),

212 uid="pandas.DataFrame",

213 source_pckg="zanj",

214 desc="pandas.DataFrame loader",

215 ),

216 # list/tuple external

217 LoaderHandler(

218 check=lambda json_item, path=None, z=None: ( # type: ignore[misc]

219 isinstance(json_item, typing.Mapping)

220 and _FORMAT_KEY in json_item

221 and json_item[_FORMAT_KEY].startswith("list")

222 and "data" in json_item

223 and isinstance(json_item["data"], typing.Sequence)

224 ),

225 load=lambda json_item, path=None, z=None: [ # type: ignore[misc]

226 load_item_recursive(x, path, z) for x in json_item["data"]

227 ],

228 uid="list",

229 source_pckg="zanj",

230 desc="list loader, for externals",

231 ),

232 LoaderHandler(

233 check=lambda json_item, path=None, z=None: ( # type: ignore[misc]

234 isinstance(json_item, typing.Mapping)

235 and _FORMAT_KEY in json_item

236 and json_item[_FORMAT_KEY].startswith("tuple")

237 and "data" in json_item

238 and isinstance(json_item["data"], typing.Sequence)

239 ),

240 load=lambda json_item, path=None, z=None: tuple( # type: ignore[misc]

241 [load_item_recursive(x, path, z) for x in json_item["data"]]

242 ),

243 uid="tuple",

244 source_pckg="zanj",

245 desc="tuple loader, for externals",

246 ),

247 ]

248}

249

250

251 def register_loader_handler(handler: LoaderHandler):

252 """register a custom loader handler"""

253 global LOADER_MAP, LOADER_MAP_LOCK

254 with LOADER_MAP_LOCK:

255 LOADER_MAP[handler.uid] = handler

256

257

258 def get_item_loader(

259 json_item: JSONitem,

260 path: ObjectPath,

261 zanj: _ZANJ_pre | None = None,

262 error_mode: ErrorMode = ErrorMode.WARN,

263 # lh_map: dict[str, LoaderHandler] = LOADER_MAP,

264) -> LoaderHandler | None:

265 """get the loader for a json item"""

266 global LOADER_MAP

267

268 # check if we recognize the format

269 if isinstance(json_item, typing.Mapping) and _FORMAT_KEY in json_item:

270 if not isinstance(json_item[_FORMAT_KEY], str):

271 raise TypeError(

272 f"invalid __muutils_format__ type '{type(json_item[_FORMAT_KEY])}' in '{path=}': '{json_item[_FORMAT_KEY] = }'"

273 )

274 if json_item[_FORMAT_KEY] in LOADER_MAP:

275 return LOADER_MAP[json_item[_FORMAT_KEY]] # type: ignore[index]

276

277 # if we dont recognize the format, try to find a loader that can handle it

278 for key, lh in LOADER_MAP.items():

279 if lh.check(json_item, path, zanj):

280 return lh

281

282 # if we still dont have a loader, return None

283 return None

284

285

286 def load_item_recursive(

287 json_item: JSONitem,

288 path: ObjectPath,

289 zanj: _ZANJ_pre | None = None,

290 error_mode: ErrorMode = ErrorMode.WARN,

291 allow_not_loading: bool = True,

292) -> Any:

293 lh: LoaderHandler | None = get_item_loader(

294 json_item=json_item,

295 path=path,

296 zanj=zanj,

297 error_mode=error_mode,

298 # lh_map=lh_map,

299 )

300

301 if lh is not None:

302 # special case for serializable dataclasses

303 if (

304 isinstance(json_item, typing.Mapping)

305 and (_FORMAT_KEY in json_item)

306 and ("SerializableDataclass" in json_item[_FORMAT_KEY]) # type: ignore[operator]

307 ):

308 # why this horribleness?

309 # SerializableDataclass, if it has a field `x` which is also a SerializableDataclass, will automatically call `x.__class__.load()`

310 # However, we need to load things in containers, as well as arrays

311 processed_json_item: dict = {

312 key: (

313 val

314 if (

315 isinstance(val, typing.Mapping)

316 and (_FORMAT_KEY in val)

317 and ("SerializableDataclass" in val[_FORMAT_KEY])

318 )

319 else load_item_recursive(

320 json_item=val,

321 path=tuple(path) + (key,),

322 zanj=zanj,

323 error_mode=error_mode,

324 )

325 )

326 for key, val in json_item.items()

327 }

328

329 return lh.load(processed_json_item, path, zanj)

330

331 else:

332 return lh.load(json_item, path, zanj)

333 else:

334 if isinstance(json_item, dict):

335 return {

336 key: load_item_recursive(

337 json_item=json_item[key],

338 path=tuple(path) + (key,),

339 zanj=zanj,

340 error_mode=error_mode,

341 # lh_map=lh_map,

342 )

343 for key in json_item

344 }

345 elif isinstance(json_item, list):

346 return [

347 load_item_recursive(

348 json_item=x,

349 path=tuple(path) + (i,),

350 zanj=zanj,

351 error_mode=error_mode,

352 # lh_map=lh_map,

353 )

354 for i, x in enumerate(json_item)

355 ]

356 elif isinstance(json_item, (str, int, float, bool, type(None))):

357 return json_item

358 else:

359 if allow_not_loading:

360 return json_item

361 else:

362 raise ValueError(

363 f"unknown type {type(json_item)} at {path}\n{json_item}"

364 )

365

366

367 def _each_item_in_externals(

368 externals: dict[str, ExternalItem],

369 json_data: JSONitem,

370) -> typing.Iterable[tuple[str, ExternalItem, Any, ObjectPath]]:

371 """note that you MUST use the raw iterator, dont try to turn into a list or something"""

372

373 sorted_externals: list[tuple[str, ExternalItem]] = sorted(

374 externals.items(), key=lambda x: len(x[1].path)

375 )

376

377 for ext_path, ext_item in sorted_externals:

378 # get the path to the item

379 path: ObjectPath = tuple(ext_item.path)

380 assert len(path) > 0

381 assert all(isinstance(key, (str, int)) for key in path), (

382 f"improper types in path {path=}"

383 )

384 # get the item

385 item = json_data

386 for i, key in enumerate(path):

387 try:

388 # ignores in this block are because we cannot know the type is indexable in static analysis

389 # but, we check the types in the line below

390 external_unloaded: bool = _populate_externals_error_checking(key, item)

391 if external_unloaded:

392 item = item["data"] # type: ignore

393 item = item[key] # type: ignore[index]

394

395 except (KeyError, IndexError, TypeError) as e:

396 raise KeyError(

397 f"could not find '{key = }' at path '{ext_path = }', specifically at index '{i = }'",

398 f"'{type(item) =}', '{len(item) = }', '{item.keys() if isinstance(item, dict) else None = }'", # type: ignore

399 f"From error: {e = }",

400 f"\n\n{item=}\n\n{ext_item=}",

401 ) from e

402

403 yield (ext_path, ext_item, item, path)

404

405

406 class LoadedZANJ:

407 """for loading a zanj file"""

408

409 def __init__(

410 self,

411 path: str | Path,

412 zanj: _ZANJ_pre,

413 ) -> None:

414 # path and zanj object

415 self._path: str = str(path)

416 self._zanj: _ZANJ_pre = zanj

417

418 # load zip file

419 _zipf: zipfile.ZipFile = zipfile.ZipFile(file=self._path, mode="r")

420

421 # load data

422 self._meta: JSONdict = json.load(_zipf.open(ZANJ_META, "r"))

423 self._json_data: JSONitem = json.load(_zipf.open(ZANJ_MAIN, "r"))

424

425 # read externals

426 self._externals: dict[str, ExternalItem] = dict()

427 for fname, ext_item in self._meta["externals_info"].items(): # type: ignore

428 item_type: str = ext_item["item_type"] # type: ignore

429 with _zipf.open(fname, "r") as fp:

430 self._externals[fname] = ExternalItem(

431 item_type=item_type, # type: ignore[arg-type]

432 data=GET_EXTERNAL_LOAD_FUNC(item_type)(self, fp),

433 path=ext_item["path"], # type: ignore

434 )

435

436 # close zip file

437 _zipf.close()

438 del _zipf

439

440 def populate_externals(self) -> None:

441 """put all external items into the main json data"""

442

443 # loop over once, populating the externals only

444 for ext_path, ext_item, item, path in _each_item_in_externals(

445 self._externals, self._json_data

446 ):

447 # replace the item with the external item

448 assert _REF_KEY in item # type: ignore

449 assert item[_REF_KEY] == ext_path # type: ignore

450 item["data"] = ext_item.data # type: ignore