add pickle module

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.
+

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)

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]);               \

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<DummyInstance>(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]));
     }));
 

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<DummyInstance>(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<DummyInstance>(t, {});
         }
 
         // __init__

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