Implement appropriate __getnewargs__ for all immutable subclassable builtin

types.  The special handling for these can now be removed from save_newobj().
Add some testing for this.

Also add support for setting the 'fast' flag on the Python Pickler class,
which suppresses use of the memo.

10 files changed, 133 insertions, 19 deletions

diff --git a/Lib/pickle.py b/Lib/pickle.py
index e36e6a6..739c24f 100644
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -191,6 +191,7 @@ class Pickler:
         self.memo = {}
         self.proto = int(proto)
         self.bin = proto >= 1
+        self.fast = 0
 
     def clear_memo(self):
         """Clears the pickler's "memo".
@@ -230,6 +231,8 @@ class Pickler:
         # But there appears no advantage to any other scheme, and this
         # scheme allows the Unpickler memo to be implemented as a plain (but
         # growable) array, indexed by memo key.
+        if self.fast:
+            return
         memo_len = len(self.memo)
         self.write(self.put(memo_len))
         self.memo[id(obj)] = memo_len, obj
@@ -378,14 +381,7 @@ class Pickler:
         if getnewargs:
             args = getnewargs()         # This bette not reference obj
         else:
-            # XXX These types should each grow a __getnewargs__
-            # implementation so this special-casing is unnecessary.
-            for cls in int, long, float, complex, str, UnicodeType, tuple:
-                if cls and isinstance(obj, cls):
-                    args = (cls(obj),)
-                    break
-            else:
-                args = ()
+            args = ()
 
         save = self.save
         write = self.write
diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
index 33b96e5..ed0e436 100644
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -324,6 +324,21 @@ class AbstractPickleTests(unittest.TestCase):
 ##         print
 ##         pickletools.dis(s)
 
+    def test_newobj_generic(self):
+        for proto in [0, 1, 2]:
+            for C in myclasses:
+                B = C.__base__
+                x = C(C.sample)
+                x.foo = 42
+                s = self.dumps(x, proto)
+##                import pickletools
+##                print
+##                pickletools.dis(s)
+                y = self.loads(s)
+                detail = (proto, C, B, x, y, type(y))
+                self.assertEqual(B(x), B(y), detail)
+                self.assertEqual(x.__dict__, y.__dict__, detail)
+
 # XXX Temporary hack, so long as the C implementation of pickle protocol
 # XXX 2 isn't ready.  When it is, move the methods in TempAbstractPickleTests
 # XXX into AbstractPickleTests above, and get rid of TempAbstractPickleTests
@@ -405,11 +420,38 @@ class TempAbstractPickleTests(unittest.TestCase):
         finally:
             copy_reg.remove_extension(__name__, "MyList", 0xfffff0)
 
+class MyInt(int):
+    sample = 1
+
+class MyLong(long):
+    sample = 1L
+
+class MyFloat(float):
+    sample = 1.0
+
+class MyComplex(complex):
+    sample = 1.0 + 0.0j
+
+class MyStr(str):
+    sample = "hello"
+
+class MyUnicode(unicode):
+    sample = u"hello \u1234"
+
 class MyTuple(tuple):
-    pass
+    sample = (1, 2, 3)
 
 class MyList(list):
-    pass
+    sample = [1, 2, 3]
+
+class MyDict(dict):
+    sample = {"a": 1, "b": 2}
+
+myclasses = [MyInt, MyLong, MyFloat,
+             # MyComplex, # XXX complex somehow doesn't work here :-(
+             MyStr, MyUnicode,
+             MyTuple, MyList, MyDict]
+
 
 class SlotList(MyList):
     __slots__ = ["foo"]
diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py
index d30e084..8870904 100644
--- a/Lib/test/test_pickle.py
+++ b/Lib/test/test_pickle.py
@@ -11,9 +11,13 @@ from test.pickletester import AbstractPersistentPicklerTests
 
 class PickleTests(AbstractPickleTests, AbstractPickleModuleTests, XXXTemp):
 
-    def setUp(self):
-        self.dumps = pickle.dumps
-        self.loads = pickle.loads
+    def dumps(self, arg, proto=0, fast=0):
+        # Ignore fast
+        return pickle.dumps(arg, proto)
+
+    def loads(self, buf):
+        # Ignore fast
+        return pickle.loads(buf)
 
     module = pickle
     error = KeyError
@@ -22,9 +26,11 @@ class PicklerTests(AbstractPickleTests):
 
     error = KeyError
 
-    def dumps(self, arg, proto=0):
+    def dumps(self, arg, proto=0, fast=0):
         f = StringIO()
         p = pickle.Pickler(f, proto)
+        if fast:
+            p.fast = fast
         p.dump(arg)
         f.seek(0)
         return f.read()
@@ -36,12 +42,14 @@ class PicklerTests(AbstractPickleTests):
 
 class PersPicklerTests(AbstractPersistentPicklerTests):
 
-    def dumps(self, arg, proto=0):
+    def dumps(self, arg, proto=0, fast=0):
         class PersPickler(pickle.Pickler):
             def persistent_id(subself, obj):
                 return self.persistent_id(obj)
         f = StringIO()
         p = PersPickler(f, proto)
+        if fast:
+            p.fast = fast
         p.dump(arg)
         f.seek(0)
         return f.read()
diff --git a/Objects/complexobject.c b/Objects/complexobject.c
index 56638d5..201da4d 100644
--- a/Objects/complexobject.c
+++ b/Objects/complexobject.c
@@ -639,8 +639,15 @@ complex_conjugate(PyObject *self)
 	return PyComplex_FromCComplex(c);
 }
 
+static PyObject *
+complex_getnewargs(PyComplexObject *v)
+{
+	return Py_BuildValue("(D)", v->cval);
+}
+
 static PyMethodDef complex_methods[] = {
 	{"conjugate",	(PyCFunction)complex_conjugate,	METH_NOARGS},
+	{"__getnewargs__",	(PyCFunction)complex_getnewargs,	METH_NOARGS},
 	{NULL,		NULL}		/* sentinel */
 };
 
diff --git a/Objects/floatobject.c b/Objects/floatobject.c
index 09406e4..6e65756 100644
--- a/Objects/floatobject.c
+++ b/Objects/floatobject.c
@@ -726,6 +726,17 @@ float_subtype_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 	return new;
 }
 
+static PyObject *
+float_getnewargs(PyFloatObject *v)
+{
+	return Py_BuildValue("(d)", v->ob_fval);
+}
+
+static PyMethodDef float_methods[] = {
+	{"__getnewargs__",	(PyCFunction)float_getnewargs,	METH_NOARGS},
+	{NULL,		NULL}		/* sentinel */
+};
+
 PyDoc_STRVAR(float_doc,
 "float(x) -> floating point number\n\
 \n\
@@ -803,7 +814,7 @@ PyTypeObject PyFloat_Type = {
 	0,					/* tp_weaklistoffset */
 	0,					/* tp_iter */
 	0,					/* tp_iternext */
-	0,					/* tp_methods */
+	float_methods,				/* tp_methods */
 	0,					/* tp_members */
 	0,					/* tp_getset */
 	0,					/* tp_base */
diff --git a/Objects/intobject.c b/Objects/intobject.c
index 805f3b7..915ef21 100644
--- a/Objects/intobject.c
+++ b/Objects/intobject.c
@@ -850,6 +850,17 @@ int_subtype_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 	return new;
 }
 
+static PyObject *
+int_getnewargs(PyIntObject *v)
+{
+	return Py_BuildValue("(l)", v->ob_ival);
+}
+
+static PyMethodDef int_methods[] = {
+	{"__getnewargs__",	(PyCFunction)int_getnewargs,	METH_NOARGS},
+	{NULL,		NULL}		/* sentinel */
+};
+
 PyDoc_STRVAR(int_doc,
 "int(x[, base]) -> integer\n\
 \n\
@@ -931,7 +942,7 @@ PyTypeObject PyInt_Type = {
 	0,					/* tp_weaklistoffset */
 	0,					/* tp_iter */
 	0,					/* tp_iternext */
-	0,					/* tp_methods */
+	int_methods,				/* tp_methods */
 	0,					/* tp_members */
 	0,					/* tp_getset */
 	0,					/* tp_base */
diff --git a/Objects/longobject.c b/Objects/longobject.c
index 1180ec2..7a04f1e 100644
--- a/Objects/longobject.c
+++ b/Objects/longobject.c
@@ -2646,6 +2646,17 @@ long_subtype_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 	return (PyObject *)new;
 }
 
+static PyObject *
+long_getnewargs(PyLongObject *v)
+{
+	return Py_BuildValue("(N)", _PyLong_Copy(v));
+}
+
+static PyMethodDef long_methods[] = {
+	{"__getnewargs__",	(PyCFunction)long_getnewargs,	METH_NOARGS},
+	{NULL,		NULL}		/* sentinel */
+};
+
 PyDoc_STRVAR(long_doc,
 "long(x[, base]) -> integer\n\
 \n\
@@ -2726,7 +2737,7 @@ PyTypeObject PyLong_Type = {
 	0,					/* tp_weaklistoffset */
 	0,					/* tp_iter */
 	0,					/* tp_iternext */
-	0,					/* tp_methods */
+	long_methods,				/* tp_methods */
 	0,					/* tp_members */
 	0,					/* tp_getset */
 	0,					/* tp_base */
diff --git a/Objects/stringobject.c b/Objects/stringobject.c
index f18edb0..9598ffb 100644
--- a/Objects/stringobject.c
+++ b/Objects/stringobject.c
@@ -3045,6 +3045,12 @@ string_splitlines(PyStringObject *self, PyObject *args)
 
 #undef SPLIT_APPEND
 
+static PyObject *
+string_getnewargs(PyStringObject *v)
+{
+	return Py_BuildValue("(s#)", v->ob_sval, v->ob_size);
+}
+
 
 static PyMethodDef
 string_methods[] = {
@@ -3091,6 +3097,7 @@ string_methods[] = {
 	 expandtabs__doc__},
 	{"splitlines", (PyCFunction)string_splitlines, METH_VARARGS,
 	 splitlines__doc__},
+	{"__getnewargs__",	(PyCFunction)string_getnewargs,	METH_NOARGS},
 	{NULL,     NULL}		     /* sentinel */
 };
 
diff --git a/Objects/tupleobject.c b/Objects/tupleobject.c
index 3a8f072..d6d0aaa 100644
--- a/Objects/tupleobject.c
+++ b/Objects/tupleobject.c
@@ -587,6 +587,18 @@ tuplesubscript(PyTupleObject* self, PyObject* item)
 	}
 }
 
+static PyObject *
+tuple_getnewargs(PyTupleObject *v)
+{
+	return Py_BuildValue("(N)", tupleslice(v, 0, v->ob_size));
+	
+}
+
+static PyMethodDef tuple_methods[] = {
+	{"__getnewargs__",	(PyCFunction)tuple_getnewargs,	METH_NOARGS},
+	{NULL,		NULL}		/* sentinel */
+};
+
 static PyMappingMethods tuple_as_mapping = {
 	(inquiry)tuplelength,
 	(binaryfunc)tuplesubscript,
@@ -625,7 +637,7 @@ PyTypeObject PyTuple_Type = {
 	0,					/* tp_weaklistoffset */
 	tuple_iter,	    			/* tp_iter */
 	0,					/* tp_iternext */
-	0,					/* tp_methods */
+	tuple_methods,				/* tp_methods */
 	0,					/* tp_members */
 	0,					/* tp_getset */
 	0,					/* tp_base */
diff --git a/Objects/unicodeobject.c b/Objects/unicodeobject.c
index 07579aa..1abef89 100644
--- a/Objects/unicodeobject.c
+++ b/Objects/unicodeobject.c
@@ -5741,6 +5741,14 @@ unicode_endswith(PyUnicodeObject *self,
 }
 
 
+
+static PyObject *
+unicode_getnewargs(PyUnicodeObject *v)
+{
+	return Py_BuildValue("(u#)", v->str, v->length);
+}
+
+
 static PyMethodDef unicode_methods[] = {
 
     /* Order is according to common usage: often used methods should
@@ -5791,6 +5799,7 @@ static PyMethodDef unicode_methods[] = {
     {"freelistsize", (PyCFunction) unicode_freelistsize, METH_NOARGS},
 #endif
 
+    {"__getnewargs__",	(PyCFunction)unicode_getnewargs, METH_NOARGS},
     {NULL, NULL}
 };