[pypy-commit] pypy numpy-pickle: Implement pickling for boxes

[rguillebert]

Author: Romain Guillebert <romain.py at gmail.com>
Branch: numpy-pickle
Changeset: r64587:ee6ad678d7a1
Date: 2013-05-27 20:17 +0200
http://bitbucket.org/pypy/pypy/changeset/ee6ad678d7a1/

Log:	Implement pickling for boxes

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -13,7 +13,7 @@
         'empty': 'interp_numarray.zeros',
         'ones': 'interp_numarray.ones',
         '_reconstruct' : 'interp_numarray._reconstruct',
-        'scalar' : 'interp_numarray.scalar',
+        'scalar' : 'interp_numarray.build_scalar',
         'dot': 'interp_arrayops.dot',
         'fromstring': 'interp_support.fromstring',
         'flatiter': 'interp_flatiter.W_FlatIterator',
diff --git a/pypy/module/micronumpy/interp_boxes.py b/pypy/module/micronumpy/interp_boxes.py
--- a/pypy/module/micronumpy/interp_boxes.py
+++ b/pypy/module/micronumpy/interp_boxes.py
@@ -11,6 +11,7 @@
 from rpython.rtyper.lltypesystem import rffi
 from rpython.tool.sourcetools import func_with_new_name
 from pypy.module.micronumpy.arrayimpl.voidbox import VoidBoxStorage
+from rpython.rlib.objectmodel import specialize
 from pypy.interpreter.mixedmodule import MixedModule
 
 MIXIN_32 = (int_typedef,) if LONG_BIT == 32 else ()
@@ -50,13 +51,24 @@
         return '%s(%s)' % (self.__class__.__name__, self.value)
 
     def descr_reduce(self, space):
+        from rpython.rlib.rstring import StringBuilder
+        from rpython.rtyper.lltypesystem import rffi, lltype
+
         numpypy = space.getbuiltinmodule("_numpypy")
         assert isinstance(numpypy, MixedModule)
         multiarray = numpypy.get("multiarray")
         assert isinstance(multiarray, MixedModule)
         scalar = multiarray.get("scalar")
 
-        return space.newtuple([scalar, space.newtuple([space.wrap(self._get_dtype(space))])])
+        value = lltype.malloc(rffi.CArray(lltype.typeOf(self.value)), 1, "raw")
+        value[0] = self.value
+
+        builder = StringBuilder()
+        builder.append_charpsize(rffi.cast(rffi.CCHARP, value), rffi.sizeof(lltype.typeOf(self.value)))
+
+        ret = space.newtuple([scalar, space.newtuple([space.wrap(self._get_dtype(space)), space.wrap(builder.build())])])
+        lltype.free(value, "raw")
+        return ret
 
 class ComplexBox(object):
     _mixin_ = True
@@ -75,13 +87,25 @@
         return dtype.box(self.imag)
 
     def descr_reduce(self, space):
+        from rpython.rlib.rstring import StringBuilder
+        from rpython.rtyper.lltypesystem import rffi, lltype
+
         numpypy = space.getbuiltinmodule("_numpypy")
         assert isinstance(numpypy, MixedModule)
         multiarray = numpypy.get("multiarray")
         assert isinstance(multiarray, MixedModule)
         scalar = multiarray.get("scalar")
 
-        return space.newtuple([scalar, space.newtuple([space.wrap(self._get_dtype(space))])])
+        value = lltype.malloc(rffi.CArray(lltype.typeOf(self.real)), 2, "raw")
+        value[0] = self.real
+        value[1] = self.imag
+
+        builder = StringBuilder()
+        builder.append_charpsize(rffi.cast(rffi.CCHARP, value), rffi.sizeof(lltype.typeOf(self.real)) * 2)
+
+        ret = space.newtuple([scalar, space.newtuple([space.wrap(self._get_dtype(space)), space.wrap(builder.build())])])
+        lltype.free(value, "raw")
+        return ret
 
 class W_GenericBox(W_Root):
     _attrs_ = ()
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
@@ -1051,8 +1051,15 @@
 def _reconstruct(space, w_subtype, w_shape, w_dtype):
     return descr_new_array(space, w_subtype, w_shape, w_dtype)
 
-def scalar(space, w_dtype):
-    pass
+def build_scalar(space, w_dtype, w_state):
+    from rpython.rtyper.lltypesystem import rffi, lltype
+
+    assert isinstance(w_dtype, interp_dtype.W_Dtype)
+
+    state = rffi.str2charp(space.str_w(w_state))
+    box = w_dtype.itemtype.box_raw_data(state)
+    lltype.free(state, "raw")
+    return box
 
 W_FlatIterator.typedef = TypeDef(
     'flatiter',
diff --git a/pypy/module/micronumpy/test/test_scalar.py b/pypy/module/micronumpy/test/test_scalar.py
--- a/pypy/module/micronumpy/test/test_scalar.py
+++ b/pypy/module/micronumpy/test/test_scalar.py
@@ -4,7 +4,7 @@
     spaceconfig = dict(usemodules=["micronumpy", "binascii", "struct"])
 
     def test_pickle(self):
-        from numpypy import dtype, int32, float64, complex128
+        from numpypy import dtype, int32, float64, complex128, zeros, sum
         from numpypy.core.multiarray import scalar
         from cPickle import loads, dumps
         i = int32(1337)
@@ -18,3 +18,6 @@
         assert loads(dumps(i)) == i
         assert loads(dumps(f)) == f
         assert loads(dumps(c)) == c
+
+        a = zeros(3)
+        assert loads(dumps(sum(a))) == sum(a)
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -145,6 +145,11 @@
         #XXX this is the place to display a warning
         return self.box(real)
 
+    def box_raw_data(self, data):
+        # For pickle
+        array = rffi.cast(rffi.CArrayPtr(self.T), data)
+        return self.box(array[0])
+
     @specialize.argtype(1)
     def unbox(self, box):
         assert isinstance(box, self.BoxType)
@@ -1108,6 +1113,11 @@
             rffi.cast(self.T, real),
             rffi.cast(self.T, imag))
 
+    def box_raw_data(self, data):
+        # For pickle
+        array = rffi.cast(rffi.CArrayPtr(self.T), data)
+        return self.box_complex(array[0], array[1])
+
     def unbox(self, box):
         assert isinstance(box, self.BoxType)
         # do this in two stages since real, imag are read only