[pypy-commit] pypy ndarray-buffer: add a new implementation of arrays which keeps alive the original buffer

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: ndarray-buffer
Changeset: r68241:193128b62e58
Date: 2013-11-19 14:50 +0100
http://bitbucket.org/pypy/pypy/changeset/193128b62e58/

Log:	add a new implementation of arrays which keeps alive the original
	buffer

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
@@ -392,7 +392,15 @@
     def __del__(self):
         free_raw_storage(self.storage, track_allocation=False)
 
+class ConcreteArrayWithBase(ConcreteArrayNotOwning):
+    def __init__(self, shape, dtype, order, strides, backstrides, storage, orig_base):
+        ConcreteArrayNotOwning.__init__(self, shape, dtype, order,
+                                        strides, backstrides, storage)
+        self.orig_base = orig_base
 
+    def base(self):
+        return self.orig_base
+        
 class NonWritableArray(ConcreteArray):
     def descr_setitem(self, space, orig_array, w_index, w_value):
         raise OperationError(space.w_ValueError, space.wrap(
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
@@ -49,11 +49,18 @@
         return W_NDimArray(impl)
 
     @staticmethod
-    def from_shape_and_storage(space, shape, storage, dtype, order='C', owning=False, w_subtype=None):
+    def from_shape_and_storage(space, shape, storage, dtype, order='C', owning=False,
+                               w_subtype=None, w_base=None):
         from pypy.module.micronumpy.arrayimpl import concrete
         assert shape
         strides, backstrides = calc_strides(shape, dtype, order)
-        if owning:
+        if w_base is not None:
+            if owning:
+                raise OperationError(space.w_ValueError, 
+                        space.wrap("Cannot have owning=True when specifying a buffer"))
+            impl = concrete.ConcreteArrayWithBase(shape, dtype, order, strides,
+                                                  backstrides, storage, w_base)
+        elif owning:
             # Will free storage when GCd
             impl = concrete.ConcreteArray(shape, dtype, order, strides,
                                                 backstrides, storage=storage)
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
@@ -1086,7 +1086,8 @@
             raise OperationError(space.w_TypeError, space.wrap(
                 "numpy scalars from buffers not supported yet"))
         storage = rffi.cast(RAW_STORAGE_PTR, raw_ptr)
-        return W_NDimArray.from_shape_and_storage(space, shape, storage, dtype)
+        return W_NDimArray.from_shape_and_storage(space, shape, storage, dtype,
+                                                  w_base=w_buffer)
 
     if not shape:
         return W_NDimArray.new_scalar(space, dtype)
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
@@ -2096,7 +2096,8 @@
         a[1] = ord('a')
         a[2] = ord('r')
         assert list(buf) == ['b', '\x00', 'a', '\x00', 'r', '\x00']
-        
+        assert a.base is buf
+
 
 class AppTestMultiDim(BaseNumpyAppTest):
     def test_init(self):