[pypy-commit] pypy numpy-pickle: Implement ndarray pickling

[rguillebert]

Author: Romain Guillebert <romain.py at gmail.com>
Branch: numpy-pickle
Changeset: r63851:08efaffe4c1f
Date: 2013-05-04 22:09 +0200
http://bitbucket.org/pypy/pypy/changeset/08efaffe4c1f/

Log:	Implement ndarray pickling

diff --git a/pypy/module/micronumpy/arrayimpl/concrete.py b/pypy/module/micronumpy/arrayimpl/concrete.py
--- a/pypy/module/micronumpy/arrayimpl/concrete.py
+++ b/pypy/module/micronumpy/arrayimpl/concrete.py
@@ -325,13 +325,14 @@
         return None
 
 class ConcreteArray(ConcreteArrayNotOwning):
-    def __init__(self, shape, dtype, order, strides, backstrides):
-        # we allocate the actual storage later because we need to compute
-        # self.size first
+    def __init__(self, shape, dtype, order, strides, backstrides, storage=None):
         null_storage = lltype.nullptr(RAW_STORAGE)
         ConcreteArrayNotOwning.__init__(self, shape, dtype, order, strides, backstrides,
                                         null_storage)
-        self.storage = dtype.itemtype.malloc(self.size)
+        if storage is None:
+            self.storage = dtype.itemtype.malloc(self.size)
+        else:
+            self.storage = storage
 
     def __del__(self):
         free_raw_storage(self.storage, track_allocation=False)
diff --git a/pypy/module/micronumpy/base.py b/pypy/module/micronumpy/base.py
--- a/pypy/module/micronumpy/base.py
+++ b/pypy/module/micronumpy/base.py
@@ -18,11 +18,11 @@
     def __init__(self, implementation):
         assert isinstance(implementation, BaseArrayImplementation)
         self.implementation = implementation
-    
+
     @staticmethod
     def from_shape(shape, dtype, order='C'):
         from pypy.module.micronumpy.arrayimpl import concrete, scalar
-        
+
         if not shape:
             impl = scalar.Scalar(dtype)
         else:
@@ -32,12 +32,17 @@
         return W_NDimArray(impl)
 
     @staticmethod
-    def from_shape_and_storage(shape, storage, dtype, order='C'):
+    def from_shape_and_storage(shape, storage, dtype, order='C', owning=False):
         from pypy.module.micronumpy.arrayimpl import concrete
         assert shape
         strides, backstrides = calc_strides(shape, dtype, order)
-        impl = concrete.ConcreteArrayNotOwning(shape, dtype, order, strides,
-                                               backstrides, storage)
+        if owning:
+            # Will free storage when GCd
+            impl = concrete.ConcreteArray(shape, dtype, order, strides,
+                                                backstrides, storage=storage)
+        else:
+            impl = concrete.ConcreteArrayNotOwning(shape, dtype, order, strides,
+                                                backstrides, storage)
         return W_NDimArray(impl)
 
 
@@ -60,7 +65,7 @@
 def convert_to_array(space, w_obj):
     from pypy.module.micronumpy.interp_numarray import array
     from pypy.module.micronumpy import interp_ufuncs
-    
+
     if isinstance(w_obj, W_NDimArray):
         return w_obj
     elif issequence_w(space, w_obj):
diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -796,6 +796,17 @@
 
         return space.newtuple([reconstruct, parameters, state])
 
+    def descr_setstate(self, space, w_state):
+        from rpython.rtyper.lltypesystem import rffi
+
+        shape = space.getitem(w_state, space.wrap(1))
+        dtype = space.getitem(w_state, space.wrap(2))
+        isfortran = space.getitem(w_state, space.wrap(3))
+        storage = space.getitem(w_state, space.wrap(4))
+
+        self.implementation = W_NDimArray.from_shape_and_storage([space.int_w(i) for i in space.listview(shape)], rffi.str2charp(space.str_w(storage), track_allocation=False), dtype, owning=True).implementation
+
+
 @unwrap_spec(offset=int)
 def descr_new_array(space, w_subtype, w_shape, w_dtype=None, w_buffer=None,
                     offset=0, w_strides=None, w_order=None):
@@ -945,7 +956,8 @@
     __pypy_data__ = GetSetProperty(W_NDimArray.fget___pypy_data__,
                                    W_NDimArray.fset___pypy_data__,
                                    W_NDimArray.fdel___pypy_data__),
-    __reduce__ = interp2app(W_NDimArray.descr_reduce)
+    __reduce__ = interp2app(W_NDimArray.descr_reduce),
+    __setstate__ = interp2app(W_NDimArray.descr_setstate),
 )
 
 @unwrap_spec(ndmin=int, copy=bool, subok=bool)
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -1749,7 +1749,7 @@
         assert data[2][4] == '\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00'
 
         pickled_data = dumps(a)
-        assert loads(pickled_data) == a
+        assert (loads(pickled_data) == a).all()
 
 class AppTestMultiDim(BaseNumpyAppTest):
     def test_init(self):