Add a version number to pickles of ndarray and dtype.

This is backwards compatible. Also added test cases to make sure.

4 files changed, 117 insertions, 29 deletions

diff --git a/numpy/core/_internal.py b/numpy/core/_internal.py
index 20117d4fd..abd0342b6 100644
--- a/numpy/core/_internal.py
+++ b/numpy/core/_internal.py
@@ -100,6 +100,10 @@ def _array_descr(descriptor):
 
     return result
 
+# Build a new array from the information in a pickle.
+# Note that the name numpy.core._internal._reconstruct is embedded in
+# the pickles of ndarrays, so don't remove the name here, or you'll
+# break backward compatibilty.
 def _reconstruct(subtype, shape, dtype):
     return ndarray.__new__(subtype, shape, dtype)
 
diff --git a/numpy/core/src/arraymethods.c b/numpy/core/src/arraymethods.c
index 79235e58c..3427ebd0b 100644
--- a/numpy/core/src/arraymethods.c
+++ b/numpy/core/src/arraymethods.c
@@ -849,6 +849,10 @@ static char doc_reduce[] = "a.__reduce__()  for pickling.";
 static PyObject *
 array_reduce(PyArrayObject *self, PyObject *args)
 {
+        /* version number of this pickle type. Increment if we need to
+           change the format. Be sure to handle the old versions in
+           array_setstate. */
+        const int version = 1;
 	PyObject *ret=NULL, *state=NULL, *obj=NULL, *mod=NULL;
 	PyObject *mybool, *thestr=NULL;
 	PyArray_Descr *descr;
@@ -873,29 +877,31 @@ array_reduce(PyArrayObject *self, PyObject *args)
 				       'b'));
 	
 	/* Now fill in object's state.  This is a tuple with 
-	   4 arguments
+	   5 arguments
 
-	   1) a Tuple giving the shape
-	   2) a PyArray_Descr Object (with correct bytorder set)
-	   3) a Bool stating if Fortran or not
-	   4) a binary string with the data (or a list for Object arrays)
+           1) an integer with the pickle version.
+	   2) a Tuple giving the shape
+	   3) a PyArray_Descr Object (with correct bytorder set)
+	   4) a Bool stating if Fortran or not
+	   5) a binary string with the data (or a list for Object arrays)
 
 	   Notice because Python does not describe a mechanism to write 
 	   raw data to the pickle, this performs a copy to a string first
 	*/
 
-	state = PyTuple_New(4);
+	state = PyTuple_New(5);
 	if (state == NULL) {
 		Py_DECREF(ret); return NULL;
 	}
-	PyTuple_SET_ITEM(state, 0, PyObject_GetAttrString((PyObject *)self, 
+        PyTuple_SET_ITEM(state, 0, PyInt_FromLong(version));
+	PyTuple_SET_ITEM(state, 1, PyObject_GetAttrString((PyObject *)self, 
 							  "shape"));
 	descr = self->descr;
 	Py_INCREF(descr);
-	PyTuple_SET_ITEM(state, 1, (PyObject *)descr);
+	PyTuple_SET_ITEM(state, 2, (PyObject *)descr);
 	mybool = (PyArray_ISFORTRAN(self) ? Py_True : Py_False);
 	Py_INCREF(mybool);
-	PyTuple_SET_ITEM(state, 2, mybool);
+	PyTuple_SET_ITEM(state, 3, mybool);
 	if (PyArray_ISOBJECT(self)) {
 		thestr = _getobject_pkl(self);
 	}
@@ -907,7 +913,7 @@ array_reduce(PyArrayObject *self, PyObject *args)
 		Py_DECREF(state);
 		return NULL;
 	}
-	PyTuple_SET_ITEM(state, 3, thestr);
+	PyTuple_SET_ITEM(state, 4, thestr);
 	PyTuple_SET_ITEM(ret, 2, state);
 	return ret;
 }
@@ -932,6 +938,7 @@ array_setstate(PyArrayObject *self, PyObject *args)
 {
 	PyObject *shape;
 	PyArray_Descr *typecode;
+        int version = 1;
 	int fortran;
 	PyObject *rawdata;
 	char *datastr;
@@ -942,10 +949,27 @@ array_setstate(PyArrayObject *self, PyObject *args)
 	/* This will free any memory associated with a and
 	   use the string in setstate as the (writeable) memory.
 	*/
-	if (!PyArg_ParseTuple(args, "(O!O!iO)", &PyTuple_Type,
+	if (!PyArg_ParseTuple(args, "(iO!O!iO)", &version, &PyTuple_Type,
+			      &shape, &PyArrayDescr_Type, &typecode,
+			      &fortran, &rawdata)) {
+            PyErr_Clear();
+            version = 0;
+	    if (!PyArg_ParseTuple(args, "(O!O!iO)", &PyTuple_Type,
 			      &shape, &PyArrayDescr_Type, &typecode, 
-			      &fortran, &rawdata))
+			      &fortran, &rawdata)) {
 		return NULL;
+            }
+        }
+
+        /* If we ever need another pickle format, increment the version
+           number. But we should still be able to handle the old versions.
+           We've only got one right now. */
+        if (version != 1 && version != 0) {
+            PyErr_Format(PyExc_ValueError,
+                         "can't handle version %d of numpy.ndarray pickle",
+                         version);
+            return NULL;
+        }
 
 	Py_XDECREF(self->descr);
 	self->descr = typecode;
diff --git a/numpy/core/src/arrayobject.c b/numpy/core/src/arrayobject.c
index ef25af71c..0b22ec510 100644
--- a/numpy/core/src/arrayobject.c
+++ b/numpy/core/src/arrayobject.c
@@ -9391,10 +9391,14 @@ arraydescr_new(PyTypeObject *subtype, PyObject *args, PyObject *kwds)
 
 static char doc_arraydescr_reduce[] = "self.__reduce__()  for pickling.";
 
-/* return a tuple of (callable object, args, state) */
+/* return a tuple of (callable object, args, state). */
 static PyObject *
 arraydescr_reduce(PyArray_Descr *self, PyObject *args)
 {
+        /* version number of this pickle type. Increment if we need to
+           change the format. Be sure to handle the old versions in
+           arraydescr_setstate. */
+        const int version = 1;
 	PyObject *ret, *mod, *obj;
 	PyObject *state;
 	char endian;
@@ -9430,15 +9434,16 @@ arraydescr_reduce(PyArray_Descr *self, PyObject *args)
 		endian = '<';
 		if (!PyArray_IsNativeByteOrder(endian)) endian = '>';
 	}
-	state = PyTuple_New(5);
-	PyTuple_SET_ITEM(state, 0, PyString_FromFormat("%c", endian));
-	PyTuple_SET_ITEM(state, 1, arraydescr_subdescr_get(self));
+	state = PyTuple_New(6);
+        PyTuple_SET_ITEM(state, 0, PyInt_FromLong(version));
+	PyTuple_SET_ITEM(state, 1, PyString_FromFormat("%c", endian));
+	PyTuple_SET_ITEM(state, 2, arraydescr_subdescr_get(self));
 	if (self->fields && self->fields != Py_None) {
 		Py_INCREF(self->fields);
-		PyTuple_SET_ITEM(state, 2, self->fields);
+		PyTuple_SET_ITEM(state, 3, self->fields);
 	}
 	else {
-		PyTuple_SET_ITEM(state, 2, Py_None);
+		PyTuple_SET_ITEM(state, 3, Py_None);
 		Py_INCREF(Py_None);
 	}
 
@@ -9449,8 +9454,8 @@ arraydescr_reduce(PyArray_Descr *self, PyObject *args)
 	}
 	else {elsize = -1; alignment = -1;}
 
-	PyTuple_SET_ITEM(state, 3, PyInt_FromLong(elsize));
-	PyTuple_SET_ITEM(state, 4, PyInt_FromLong(alignment));
+	PyTuple_SET_ITEM(state, 4, PyInt_FromLong(elsize));
+	PyTuple_SET_ITEM(state, 5, PyInt_FromLong(alignment));
 
 	PyTuple_SET_ITEM(ret, 2, state);
 	return ret;
@@ -9465,15 +9470,33 @@ static PyObject *
 arraydescr_setstate(PyArray_Descr *self, PyObject *args)
 {
 	int elsize = -1, alignment = -1;
+        int version = 1;
 	char endian;
 	PyObject *subarray, *fields;
 
 	if (self->fields == Py_None) {Py_INCREF(Py_None); return Py_None;}
 
-	if (!PyArg_ParseTuple(args, "(cOOii)", &endian, &subarray, &fields,
-			      &elsize, &alignment)) return NULL;
+	if (!PyArg_ParseTuple(args, "(icOOii)", &version, &endian, &subarray,
+                              &fields, &elsize, &alignment)) {
+            PyErr_Clear();
+            version = 0;
+	    if (!PyArg_ParseTuple(args, "(cOOii)", &endian, &subarray,
+                                  &fields, &elsize, &alignment)) {
+                return NULL;
+            }
+        }
+
+        /* If we ever need another pickle format, increment the version
+           number. But we should still be able to handle the old versions.
+           We've only got one right now. */
+        if (version != 1 && version != 0) {
+            PyErr_Format(PyExc_ValueError,
+                         "can't handle version %d of numpy.dtype pickle",
+                          version);
+            return NULL;
+        }
 
-	if (endian != '|' && 
+	if (endian != '|' &&
 	    PyArray_IsNativeByteOrder(endian)) endian = '=';
 
 	self->byteorder = endian;
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index f7251d4aa..a0c7db466 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -237,14 +237,51 @@ class test_subscripting(ScipyTestCase):
         self.failUnless(type(x[0, ...]) is ndarray)
 
 class test_pickling(ScipyTestCase):
-    def setUp(self):
-        self.carray = array([[2,9],[7,0],[3,8]])
-        self.tarray = transpose(self.carray)
-
     def check_both(self):
         import pickle
-        assert_equal(self.carray, pickle.loads(self.carray.dumps()))
-        assert_equal(self.tarray, pickle.loads(self.tarray.dumps()))
+        carray = array([[2,9],[7,0],[3,8]])
+        tarray = transpose(carray)
+        assert_equal(carray, pickle.loads(carray.dumps()))
+        assert_equal(tarray, pickle.loads(tarray.dumps()))
+
+    # version 0 pickles, using protocol=2 to pickle
+    # version 0 doesn't have a version field
+    def check_version0_int8(self):
+        s = '\x80\x02cnumpy.core._internal\n_reconstruct\nq\x01cnumpy\nndarray\nq\x02K\x00\x85U\x01b\x87Rq\x03(K\x04\x85cnumpy\ndtype\nq\x04U\x02i1K\x00K\x01\x87Rq\x05(U\x01|NNJ\xff\xff\xff\xffJ\xff\xff\xff\xfftb\x89U\x04\x01\x02\x03\x04tb.'
+        a = array([1,2,3,4], dtype=int8)
+        p = loads(s)
+        assert_equal(a, p)
+
+    def check_version0_float32(self):
+        s = '\x80\x02cnumpy.core._internal\n_reconstruct\nq\x01cnumpy\nndarray\nq\x02K\x00\x85U\x01b\x87Rq\x03(K\x04\x85cnumpy\ndtype\nq\x04U\x02f4K\x00K\x01\x87Rq\x05(U\x01<NNJ\xff\xff\xff\xffJ\xff\xff\xff\xfftb\x89U\x10\x00\x00\x80?\x00\x00\x00@\x00\x00@@\x00\x00\x80@tb.'
+        a = array([1.0, 2.0, 3.0, 4.0], dtype=float32)
+        p = loads(s)
+        assert_equal(a, p)
+
+    def check_version0_object(self):
+        s = '\x80\x02cnumpy.core._internal\n_reconstruct\nq\x01cnumpy\nndarray\nq\x02K\x00\x85U\x01b\x87Rq\x03(K\x02\x85cnumpy\ndtype\nq\x04U\x02O8K\x00K\x01\x87Rq\x05(U\x01|NNJ\xff\xff\xff\xffJ\xff\xff\xff\xfftb\x89]q\x06(}q\x07U\x01aK\x01s}q\x08U\x01bK\x02setb.'
+        a = array([{'a':1}, {'b':2}])
+        p = loads(s)
+        assert_equal(a, p)
+
+    # version 1 pickles, using protocol=2 to pickle
+    def check_version1_int8(self):
+        s = '\x80\x02cnumpy.core._internal\n_reconstruct\nq\x01cnumpy\nndarray\nq\x02K\x00\x85U\x01b\x87Rq\x03(K\x01K\x04\x85cnumpy\ndtype\nq\x04U\x02i1K\x00K\x01\x87Rq\x05(K\x01U\x01|NNJ\xff\xff\xff\xffJ\xff\xff\xff\xfftb\x89U\x04\x01\x02\x03\x04tb.'
+        a = array([1,2,3,4], dtype=int8)
+        p = loads(s)
+        assert_equal(a, p)
+
+    def check_version1_float32(self):
+        s = '\x80\x02cnumpy.core._internal\n_reconstruct\nq\x01cnumpy\nndarray\nq\x02K\x00\x85U\x01b\x87Rq\x03(K\x01K\x04\x85cnumpy\ndtype\nq\x04U\x02f4K\x00K\x01\x87Rq\x05(K\x01U\x01<NNJ\xff\xff\xff\xffJ\xff\xff\xff\xfftb\x89U\x10\x00\x00\x80?\x00\x00\x00@\x00\x00@@\x00\x00\x80@tb.'
+        a = array([1.0, 2.0, 3.0, 4.0], dtype=float32)
+        p = loads(s)
+        assert_equal(a, p)
+
+    def check_version1_object(self):
+        s = '\x80\x02cnumpy.core._internal\n_reconstruct\nq\x01cnumpy\nndarray\nq\x02K\x00\x85U\x01b\x87Rq\x03(K\x01K\x02\x85cnumpy\ndtype\nq\x04U\x02O8K\x00K\x01\x87Rq\x05(K\x01U\x01|NNJ\xff\xff\xff\xffJ\xff\xff\xff\xfftb\x89]q\x06(}q\x07U\x01aK\x01s}q\x08U\x01bK\x02setb.'
+        a = array([{'a':1}, {'b':2}])
+        p = loads(s)
+        assert_equal(a, p)
 
 class test_fancy_indexing(ScipyTestCase): 
     def check_list(self):