diff --git a/docs/modules/pickle.md b/docs/modules/pickle.md new file mode 100644 index 00000000..64474bb3 --- /dev/null +++ b/docs/modules/pickle.md @@ -0,0 +1,13 @@ +--- +icon: package +label: pickle +--- + +### `pickle.dumps(obj) -> bytes` + +Return the pickled representation of an object as a bytes object. + +### `pickle.loads(b: bytes)` + +Return the unpickled object from a bytes object. + diff --git a/python/pickle.py b/python/pickle.py new file mode 100644 index 00000000..3a3a65f2 --- /dev/null +++ b/python/pickle.py @@ -0,0 +1,90 @@ +import json +import builtins + +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__ + raise PickleError(f"cannot find __new__ for {cls.__name__}") + +def _wrap(o): + if type(o) in (int, float, str, bool, type(None)): + return o + if type(o) is list: + return ["list", [_wrap(i) for i in o]] + if type(o) is tuple: + return ["tuple", [_wrap(i) for i in o]] + if type(o) is dict: + return ["dict", [[_wrap(k), _wrap(v)] for k,v in o.items()]] + if type(o) is bytes: + return ["bytes", [o[j] for j in range(len(o))]] + + _0 = o.__class__.__name__ + if hasattr(o, "__getnewargs__"): + _1 = o.__getnewargs__() # an iterable + _1 = [_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] = _wrap(v) + return [_0, _1, _2] + +def _unwrap(o): + if type(o) in (int, float, str, bool, type(None)): + return o + if isinstance(o, list): + if o[0] == "list": + return [_unwrap(i) for i in o[1]] + if o[0] == "tuple": + return tuple([_unwrap(i) for i in o[1]]) + if o[0] == "dict": + return {_unwrap(k): _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 = [_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, _unwrap(v)) + return inst + raise PickleError(f"cannot unpickle {type(o).__name__} object") + + +def dumps(o) -> bytes: + return json.dumps(_wrap(o)).encode() + + +def loads(b) -> object: + assert type(b) is bytes + o = json.loads(b.decode()) + return _unwrap(o) \ No newline at end of file diff --git a/src/linalg.h b/src/linalg.h index 6a45adae..00348b7a 100644 --- a/src/linalg.h +++ b/src/linalg.h @@ -330,6 +330,11 @@ struct PyVec2: Vec2 { return VAR(Vec2(x, y)); }); + vm->bind_method<0>(type, "__getnewargs__", [](VM* vm, ArgsView args){ + PyVec2& self = _CAST(PyVec2&, args[0]); + return VAR(Tuple({ VAR(self.x), VAR(self.y) })); + }); + vm->bind__repr__(OBJ_GET(Type, type), [](VM* vm, PyObject* obj){ PyVec2& self = _CAST(PyVec2&, obj); std::stringstream ss; @@ -384,6 +389,11 @@ struct PyVec3: Vec3 { return VAR(Vec3(x, y, z)); }); + vm->bind_method<0>(type, "__getnewargs__", [](VM* vm, ArgsView args){ + PyVec3& self = _CAST(PyVec3&, args[0]); + return VAR(Tuple({ VAR(self.x), VAR(self.y), VAR(self.z) })); + }); + vm->bind__repr__(OBJ_GET(Type, type), [](VM* vm, PyObject* obj){ PyVec3& self = _CAST(PyVec3&, obj); std::stringstream ss; @@ -444,6 +454,13 @@ struct PyMat3x3: Mat3x3{ return vm->None; }); + vm->bind_method<0>(type, "__getnewargs__", [](VM* vm, ArgsView args){ + PyMat3x3& self = _CAST(PyMat3x3&, args[0]); + Tuple t(9); + for(int i=0; i<9; i++) t[i] = VAR(self.v[i]); + return VAR(std::move(t)); + }); + #define METHOD_PROXY_NONE(name) \ vm->bind_method<0>(type, #name, [](VM* vm, ArgsView args){ \ PyMat3x3& self = _CAST(PyMat3x3&, args[0]); \ diff --git a/src/pocketpy.h b/src/pocketpy.h index bc62198b..85575e0a 100644 --- a/src/pocketpy.h +++ b/src/pocketpy.h @@ -219,6 +219,12 @@ inline void init_builtins(VM* _vm) { _vm->bind__eq__(_vm->tp_object, [](VM* vm, PyObject* lhs, PyObject* rhs) { return lhs == rhs; }); _vm->bind__hash__(_vm->tp_object, [](VM* vm, PyObject* obj) { return BITS(obj); }); + _vm->cached_object__new__ = _vm->bind_constructor<1>("object", [](VM* vm, ArgsView args) { + vm->check_non_tagged_type(args[0], vm->tp_type); + Type t = OBJ_GET(Type, args[0]); + return vm->heap.gcnew(t, {}); + }); + _vm->bind_constructor<2>("type", CPP_LAMBDA(vm->_t(args[1]))); _vm->bind_constructor<-1>("range", [](VM* vm, ArgsView args) { @@ -1285,7 +1291,7 @@ inline void VM::post_init(){ add_module_random(this); add_module_base64(this); - for(const char* name: {"this", "functools", "collections", "heapq", "bisect"}){ + for(const char* name: {"this", "functools", "collections", "heapq", "bisect", "pickle"}){ _lazy_modules[name] = kPythonLibs[name]; } @@ -1327,7 +1333,7 @@ inline void VM::post_init(){ })); _t(tp_object)->attr().set("__dict__", property([](VM* vm, ArgsView args){ - if(is_tagged(args[0]) || !args[0]->is_attr_valid()) vm->AttributeError("'__dict__'"); + if(is_tagged(args[0]) || !args[0]->is_attr_valid()) return vm->None; return VAR(MappingProxy(args[0])); })); diff --git a/src/vm.h b/src/vm.h index f50d7dd0..dba98671 100644 --- a/src/vm.h +++ b/src/vm.h @@ -133,6 +133,8 @@ public: Type tp_super, tp_exception, tp_bytes, tp_mappingproxy; Type tp_dict, tp_property, tp_star_wrapper; + PyObject* cached_object__new__; + const bool enable_os; VM(bool enable_os=true) : heap(this), enable_os(enable_os) { @@ -1323,7 +1325,14 @@ inline PyObject* VM::vectorcall(int ARGC, int KWARGC, bool op_call){ DEF_SNAME(__new__); PyObject* new_f = find_name_in_mro(callable, __new__); PyObject* obj; - if(new_f != nullptr){ +#if DEBUG_EXTRA_CHECK + PK_ASSERT(new_f != nullptr); +#endif + if(new_f == cached_object__new__) { + // fast path for object.__new__ + Type t = OBJ_GET(Type, callable); + obj= vm->heap.gcnew(t, {}); + }else{ PUSH(new_f); PUSH(PY_NULL); PUSH(callable); // cls @@ -1331,10 +1340,6 @@ inline PyObject* VM::vectorcall(int ARGC, int KWARGC, bool op_call){ for(PyObject* obj: kwargs) PUSH(obj); // if obj is not an instance of callable, the behavior is undefined obj = vectorcall(ARGC+1, KWARGC); - }else{ - // fast path for object.__new__ - Type t = OBJ_GET(Type, callable); - obj= vm->heap.gcnew(t, {}); } // __init__ diff --git a/tests/81_pickle.py b/tests/81_pickle.py new file mode 100644 index 00000000..baba8783 --- /dev/null +++ b/tests/81_pickle.py @@ -0,0 +1,45 @@ +from pickle import dumps, loads, _wrap, _unwrap + +def test(x, y): + _0 = _wrap(x) + _1 = _unwrap(y) + assert _0 == y, f"{_0} != {y}" + assert _1 == x, f"{_1} != {x}" + assert x == loads(dumps(x)) + +test(1, 1) +test(1.0, 1.0) +test("hello", "hello") +test(True, True) +test(False, False) +test(None, None) + +test([1, 2, 3], ["list", [1, 2, 3]]) +test((1, 2, 3), ["tuple", [1, 2, 3]]) +test({1: 2, 3: 4}, ["dict", [[1, 2], [3, 4]]]) + +class Foo: + def __init__(self, x, y): + self.x = x + self.y = y + + def __eq__(self, __value: object) -> bool: + if not isinstance(__value, Foo): + return False + return self.x == __value.x and self.y == __value.y + + def __repr__(self) -> str: + return f"Foo({self.x}, {self.y})" + +foo = Foo(1, 2) +test(foo, ["__main__.Foo", None, {"x": 1, "y": 2}]) + +from linalg import vec2 + +test(vec2(1, 2), ["linalg.vec2", [1, 2], None]) + +a = {1, 2, 3, 4} +test(a, ['set', None, {'_a': ['dict', [[1, None], [2, None], [3, None], [4, None]]]}]) + +a = bytes([1, 2, 3, 4]) +assert loads(dumps(a)) == a \ No newline at end of file