Improve pickling of builtins

src/runtime/float.cpp

@@ -1626,6 +1626,10 @@ exit:
     return result;
 }
 
+static PyObject* float_getnewargs(PyFloatObject* v) noexcept {
+    return Py_BuildValue("(d)", v->ob_fval);
+}
+
 static PyMethodDef float_methods[] = { { "hex", (PyCFunction)float_hex, METH_NOARGS, NULL },
                                        { "fromhex", (PyCFunction)float_fromhex, METH_O | METH_CLASS, NULL },
                                        { "as_integer_ratio", (PyCFunction)float_as_integer_ratio, METH_NOARGS, NULL },
@@ -1637,6 +1641,9 @@ void setupFloat() {
     static PyNumberMethods float_as_number;
     float_cls->tp_as_number = &float_as_number;
 
+    float_cls->giveAttr("__getnewargs__", new BoxedFunction(boxRTFunction((void*)float_getnewargs, UNKNOWN, 1,
+                                                                          ParamNames::empty(), CAPI)));
+
     _addFunc("__add__", BOXED_FLOAT, (void*)floatAddFloat, (void*)floatAddInt, (void*)floatAdd);
     float_cls->giveAttr("__radd__", float_cls->getattr(internStringMortal("__add__")));
 

src/runtime/int.cpp

@@ -1206,6 +1206,10 @@ static PyObject* int_richcompare(PyObject* v, PyObject* w, int op) noexcept {
     }
 }
 
+static PyObject* int_getnewargs(BoxedInt* v) noexcept {
+    return Py_BuildValue("(l)", v->n);
+}
+
 void setupInt() {
     static PyNumberMethods int_as_number;
     int_cls->tp_as_number = &int_as_number;
@@ -1215,6 +1219,9 @@ void setupInt() {
         gc::registerPermanentRoot(interned_ints[i]);
     }
 
+    int_cls->giveAttr("__getnewargs__",
+                      new BoxedFunction(boxRTFunction((void*)int_getnewargs, UNKNOWN, 1, ParamNames::empty(), CAPI)));
+
     _addFuncIntFloatUnknown("__add__", (void*)intAddInt, (void*)intAddFloat, (void*)intAdd);
     _addFuncIntUnknown("__and__", BOXED_INT, (void*)intAndInt, (void*)intAnd);
     _addFuncIntUnknown("__or__", BOXED_INT, (void*)intOrInt, (void*)intOr);

src/runtime/long.cpp

@@ -74,7 +74,7 @@ extern "C" int _PyLong_Sign(PyObject* l) noexcept {
 
 extern "C" PyObject* _PyLong_Copy(PyLongObject* src) noexcept {
     BoxedLong* rtn = new BoxedLong();
-    mpz_init_set(((BoxedLong*)src)->n, rtn->n);
+    mpz_init_set(rtn->n, ((BoxedLong*)src)->n);
     return rtn;
 }
 
@@ -1449,6 +1449,10 @@ static Box* long1(Box* b, void*) {
     return boxLong(1);
 }
 
+static PyObject* long_getnewargs(PyLongObject* v) noexcept {
+    return Py_BuildValue("(N)", _PyLong_Copy(v));
+}
+
 void setupLong() {
     static PyNumberMethods long_as_number;
     long_cls->tp_as_number = &long_as_number;
@@ -1515,6 +1519,9 @@ void setupLong() {
     long_cls->giveAttr("numerator", new (pyston_getset_cls) BoxedGetsetDescriptor(longLong, NULL, NULL));
     long_cls->giveAttr("denominator", new (pyston_getset_cls) BoxedGetsetDescriptor(long1, NULL, NULL));
 
+    long_cls->giveAttr("__getnewargs__",
+                       new BoxedFunction(boxRTFunction((void*)long_getnewargs, UNKNOWN, 1, ParamNames::empty(), CAPI)));
+
     add_operators(long_cls);
     long_cls->freeze();
 

src/runtime/long.h

@@ -41,6 +41,8 @@ public:
 extern "C" Box* createLong(llvm::StringRef s);
 extern "C" BoxedLong* boxLong(int64_t n);
 
+Box* longRepr(BoxedLong* v);
+
 Box* longNeg(BoxedLong* lhs);
 Box* longAbs(BoxedLong* v1);
 

src/runtime/set.cpp

@@ -637,6 +637,30 @@ extern "C" PyObject* PyFrozenSet_New(PyObject* iterable) noexcept {
     }
 }
 
+static PyObject* set_reduce(BoxedSet* so) noexcept {
+    PyObject* keys = NULL, * args = NULL, * result = NULL, * dict = NULL;
+
+    keys = PySequence_List((PyObject*)so);
+    if (keys == NULL)
+        goto done;
+    args = PyTuple_Pack(1, keys);
+    if (args == NULL)
+        goto done;
+    dict = PyObject_GetAttrString((PyObject*)so, "__dict__");
+    if (dict == NULL) {
+        PyErr_Clear();
+        dict = Py_None;
+        Py_INCREF(dict);
+    }
+    result = PyTuple_Pack(3, Py_TYPE(so), args, dict);
+done:
+    Py_XDECREF(args);
+    Py_XDECREF(keys);
+    Py_XDECREF(dict);
+    return result;
+}
+
+
 } // namespace set
 
 void BoxedSet::dealloc(Box* b) noexcept {
@@ -646,6 +670,13 @@ void BoxedSet::dealloc(Box* b) noexcept {
 
 using namespace pyston::set;
 
+static PyMethodDef set_methods[] = {
+    { "__reduce__", (PyCFunction)set_reduce, METH_NOARGS, NULL },
+};
+static PyMethodDef frozenset_methods[] = {
+    { "__reduce__", (PyCFunction)set_reduce, METH_NOARGS, NULL },
+};
+
 void setupSet() {
     static PySequenceMethods set_as_sequence;
     set_cls->tp_as_sequence = &set_as_sequence;
@@ -780,6 +811,14 @@ void setupSet() {
     set_cls->giveAttr("copy", new BoxedFunction(boxRTFunction((void*)setCopy, UNKNOWN, 1)));
     set_cls->giveAttr("pop", new BoxedFunction(boxRTFunction((void*)setPop, UNKNOWN, 1)));
 
+    for (auto& md : set_methods) {
+        set_cls->giveAttr(md.ml_name, new BoxedMethodDescriptor(&md, set_cls));
+    }
+
+    for (auto& md : frozenset_methods) {
+        frozenset_cls->giveAttr(md.ml_name, new BoxedMethodDescriptor(&md, frozenset_cls));
+    }
+
     set_cls->freeze();
     frozenset_cls->freeze();
 

src/runtime/str.cpp

@@ -2786,6 +2786,10 @@ static int string_buffer_getbuffer(BoxedString* self, Py_buffer* view, int flags
     return PyBuffer_FillInfo(view, (PyObject*)self, self->data(), self->size(), 1, flags);
 }
 
+static PyObject* string_getnewargs(BoxedString* v) noexcept {
+    return Py_BuildValue("(s#)", v->data(), v->size());
+}
+
 static PyBufferProcs string_as_buffer = {
     (readbufferproc)string_buffer_getreadbuf, // comments are the only way I've found of
     (writebufferproc)NULL,                    // forcing clang-format to break these onto multiple lines
@@ -2834,6 +2838,9 @@ void setupStr() {
     str_cls->tp_as_buffer = &string_as_buffer;
     str_cls->tp_print = string_print;
 
+    str_cls->giveAttr("__getnewargs__", new BoxedFunction(boxRTFunction((void*)string_getnewargs, UNKNOWN, 1,
+                                                                        ParamNames::empty(), CAPI)));
+
     str_cls->giveAttr("__len__", new BoxedFunction(boxRTFunction((void*)strLen, BOXED_INT, 1)));
     str_cls->giveAttr("__str__", new BoxedFunction(boxRTFunction((void*)strStr, STR, 1)));
     str_cls->giveAttr("__repr__", new BoxedFunction(boxRTFunction((void*)strRepr, STR, 1)));

src/runtime/tuple.cpp

@@ -643,6 +643,12 @@ extern "C" void BoxedTupleIterator::gcHandler(GCVisitor* v, Box* b) {
     v->visit(&it->t);
 }
 
+static Box* tuple_getnewargs(Box* _self) noexcept {
+    RELEASE_ASSERT(PyTuple_Check(_self), "");
+    PyTupleObject* v = reinterpret_cast<PyTupleObject*>(_self);
+    return Py_BuildValue("(N)", tupleslice(v, 0, Py_SIZE(v)));
+}
+
 void setupTuple() {
     static PySequenceMethods tuple_as_sequence;
     tuple_cls->tp_as_sequence = &tuple_as_sequence;
@@ -681,6 +687,8 @@ void setupTuple() {
     tuple_cls->giveAttr("__mul__", new BoxedFunction(boxRTFunction((void*)tupleMul, BOXED_TUPLE, 2)));
     tuple_cls->giveAttr("__rmul__", new BoxedFunction(boxRTFunction((void*)tupleMul, BOXED_TUPLE, 2)));
 
+    tuple_cls->giveAttr("__getnewargs__", new BoxedFunction(boxRTFunction((void*)tuple_getnewargs, UNKNOWN, 1,
+                                                                          ParamNames::empty(), CAPI)));
 
     tuple_cls->tp_hash = (hashfunc)tuple_hash;
     tuple_cls->tp_as_sequence->sq_slice = (ssizessizeargfunc)&tupleslice;

src/runtime/util.cpp

@@ -17,6 +17,7 @@
 #include "codegen/codegen.h"
 #include "core/options.h"
 #include "core/types.h"
+#include "runtime/long.h"
 #include "runtime/objmodel.h"
 
 // Temp hack to get CType sort of importing
@@ -238,6 +239,11 @@ extern "C" void dumpEx(void* p, int levels) {
             printf("Int value: %ld\n", static_cast<BoxedInt*>(b)->n);
         }
 
+        if (PyLong_Check(b)) {
+            PyObject* str = longRepr(static_cast<BoxedLong*>(b));
+            printf("Long value: %s\n", static_cast<BoxedString*>(str)->c_str());
+        }
+
         if (PyList_Check(b)) {
             auto l = static_cast<BoxedList*>(b);
             printf("%ld elements\n", l->size);

test/tests/pickle_test.py

@@ -22,3 +22,23 @@ c = C()
 c.a = 1
 print repr(pickle.dumps(c))
 print pickle.loads(pickle.dumps(c)).a
+
+
+for obj in [(1, 2), "hello world", u"hola world", 1.0, 1j, 1L, 2, {1:2}, set([3]), frozenset([4])]:
+    print
+    print "Testing pickling subclasses of %s..." % type(obj)
+
+    class MySubclass(type(obj)):
+        pass
+
+    o = MySubclass(obj)
+    print repr(o), type(o)
+
+    for protocol in (0, 1, 2):
+        print "Protocol", protocol
+
+        s = pickle.dumps(o, protocol=protocol)
+        print repr(s)
+
+        o2 = pickle.loads(s)
+        print repr(o2), type(o2)