import json import builtins _BASIC_TYPES = [int, float, str, bool, type(None)] def _find_class(path: str): if "." not in path: g = globals() if path in g: return g[path] return builtins.__dict__[path] modname, name = path.split(".") return __import__(modname).__dict__[name] def _find__new__(cls): while cls is not None: d = cls.__dict__ if "__new__" in d: return d["__new__"] cls = cls.__base__ assert False class _Pickler: def __init__(self) -> None: self.raw_memo = {} # id -> int self.memo = [] # int -> object def wrap(self, o): if type(o) in _BASIC_TYPES: return o index = self.raw_memo.get(id(o), None) if index is not None: return ["$", index] ret = [] index = len(self.memo) self.memo.append(ret) self.raw_memo[id(o)] = index if type(o) is list: ret.append("list") ret.append([self.wrap(i) for i in o]) return ["$", index] if type(o) is tuple: ret.append("tuple") ret.append([self.wrap(i) for i in o]) return ["$", index] if type(o) is dict: ret.append("dict") ret.append([[self.wrap(k), self.wrap(v)] for k,v in o.items()]) return ["$", index] if type(o) is bytes: ret.append("bytes") ret.append([o[j] for j in range(len(o))]) return ["$", index] _0 = o.__class__.__name__ if hasattr(o, "__getnewargs__"): _1 = o.__getnewargs__() # an iterable _1 = [self.wrap(i) for i in _1] else: _1 = None if hasattr(o, "__getstate__"): _2 = o.__getstate__() else: if o.__dict__ is None: _2 = None else: _2 = {} for k,v in o.__dict__.items(): _2[k] = self.wrap(v) ret.append(_0) ret.append(_1) ret.append(_2) return ["$", index] class _Unpickler: def __init__(self, memo: list) -> None: self.memo = memo self.unwrapped = [None] * len(memo) def unwrap_ref(self, i: int): if self.unwrapped[i] is None: o = self.memo[i] assert type(o) is list assert o[0] != '$' self.unwrapped[i] = self.unwrap(o) return self.unwrapped[i] def unwrap(self, o): if type(o) in _BASIC_TYPES: return o assert type(o) is list if o[0] == '$': index = o[1] return self.unwrap_ref(index) if o[0] == "list": return [self.unwrap(i) for i in o[1]] if o[0] == "tuple": return tuple([self.unwrap(i) for i in o[1]]) if o[0] == "dict": return {self.unwrap(k): self.unwrap(v) for k,v in o[1]} if o[0] == "bytes": return bytes(o[1]) # generic object cls, newargs, state = o cls = _find_class(o[0]) # create uninitialized instance new_f = _find__new__(cls) if newargs is not None: newargs = [self.unwrap(i) for i in newargs] inst = new_f(cls, *newargs) else: inst = new_f(cls) # restore state if hasattr(inst, "__setstate__"): inst.__setstate__(state) else: if state is not None: for k,v in state.items(): setattr(inst, k, self.unwrap(v)) return inst def _wrap(o): p = _Pickler() o = p.wrap(o) return [o, p.memo] def _unwrap(packed: list): o, memo = packed return _Unpickler(memo).unwrap(o) def dumps(o) -> bytes: o = _wrap(o) return json.dumps(o).encode() def loads(b) -> object: assert type(b) is bytes o = json.loads(b.decode()) return _unwrap(o)