Commits

Antonio Cuni committed c437c02

add a way to construct a numpy array from an existing buffer. PyPy-only implementation detail

Comments (0)

Files changed (4)

pypy/module/micronumpy/arrayimpl/concrete.py

 from pypy.interpreter.error import OperationError, operationerrfmt
 from pypy.rpython.lltypesystem import rffi, lltype
 from pypy.rlib import jit
-from pypy.rlib.rawstorage import free_raw_storage
+from pypy.rlib.rawstorage import free_raw_storage, RAW_STORAGE
 from pypy.rlib.debug import make_sure_not_resized
 
 class ConcreteArrayIterator(base.BaseArrayIterator):
     def get_storage_as_int(self, space):
         return rffi.cast(lltype.Signed, self.storage)
 
-class ConcreteArray(BaseConcreteArray):
-    def __init__(self, shape, dtype, order, strides, backstrides):
+class ConcreteArrayNotOwning(BaseConcreteArray):
+    def __init__(self, shape, dtype, order, strides, backstrides, storage):
         make_sure_not_resized(shape)
         make_sure_not_resized(strides)
         make_sure_not_resized(backstrides)
         self.shape = shape
         self.size = support.product(shape) * dtype.get_size()
-        self.storage = dtype.itemtype.malloc(self.size)
         self.order = order
         self.dtype = dtype
         self.strides = strides
         self.backstrides = backstrides
+        self.storage = storage
 
     def create_iter(self, shape=None):
         if shape is None or shape == self.get_shape():
     def fill(self, box):
         self.dtype.fill(self.storage, box, 0, self.size)
 
-    def __del__(self):
-        free_raw_storage(self.storage, track_allocation=False)
-
     def set_shape(self, space, new_shape):
         strides, backstrides = support.calc_strides(new_shape, self.dtype,
                                                     self.order)
         return SliceArray(0, strides, backstrides, new_shape, self)
 
+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
+        null_storage = lltype.nullptr(RAW_STORAGE)
+        ConcreteArrayNotOwning.__init__(self, shape, dtype, order, strides, backstrides,
+                                        null_storage)
+        self.storage = dtype.itemtype.malloc(self.size)
+
+    def __del__(self):
+        free_raw_storage(self.storage, track_allocation=False)
+
+
 class NonWritableArray(ConcreteArray):
     def descr_setitem(self, space, w_index, w_value):
         raise OperationError(space.w_RuntimeError, space.wrap(

pypy/module/micronumpy/base.py

         return W_NDimArray(impl)
 
     @staticmethod
+    def from_shape_and_storage(shape, storage, dtype, order='C'):
+        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)
+        return W_NDimArray(impl)
+
+
+    @staticmethod
     def new_slice(offset, strides, backstrides, shape, parent, dtype=None):
         from pypy.module.micronumpy.arrayimpl import concrete
 

pypy/module/micronumpy/interp_numarray.py

                                                        space.w_False]))
         return w_d
 
-
     # --------------------- operations ----------------------------
 
     def _unaryop_impl(ufunc_name):
         return W_NDimArray.new_scalar(space, dtype)
     return W_NDimArray.from_shape(shape, dtype)
 
+@unwrap_spec(addr=int)
+def descr__from_shape_and_storage(space, w_cls, w_shape, addr, w_dtype):
+    """
+    Create an array from an existing buffer, given its address as int.
+    PyPy-only implementation detail.
+    """
+    from pypy.rpython.lltypesystem import rffi
+    from pypy.rlib.rawstorage import RAW_STORAGE_PTR
+    shape = _find_shape(space, w_shape)
+    storage = rffi.cast(RAW_STORAGE_PTR, addr)
+    dtype = space.interp_w(interp_dtype.W_Dtype,
+                           space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype))
+    return W_NDimArray.from_shape_and_storage(shape, storage, dtype)
+
 W_NDimArray.typedef = TypeDef(
     "ndarray",
     __new__ = interp2app(descr_new_array),
     imag = GetSetProperty(W_NDimArray.descr_get_imag,
                           W_NDimArray.descr_set_imag),
     __array_interface__ = GetSetProperty(W_NDimArray.descr_array_iface),
+   _from_shape_and_storage = interp2app(descr__from_shape_and_storage,
+                                        as_classmethod=True)
 )
 
 @unwrap_spec(ndmin=int, copy=bool, subok=bool)

pypy/module/micronumpy/test/test_numarray.py

         assert shape == [2]
         assert space.str_w(elems[0]) == "a"
         assert space.str_w(elems[1]) == "b"
-        
+
+    def test_from_shape_and_storage(self):
+        from pypy.rlib.rawstorage import alloc_raw_storage, raw_storage_setitem
+        from pypy.rpython.lltypesystem import rffi
+        from pypy.module.micronumpy.interp_dtype import get_dtype_cache
+        storage = alloc_raw_storage(4, track_allocation=False, zero=True)
+        for i in range(4):
+            raw_storage_setitem(storage, i, rffi.cast(rffi.UCHAR, i))
+        #
+        dtypes = get_dtype_cache(self.space)
+        w_array = W_NDimArray.from_shape_and_storage([2, 2], storage, dtypes.w_int8dtype)
+        def get(i, j):
+            return w_array.getitem(self.space, [i, j]).value
+        assert get(0, 0) == 0
+        assert get(0, 1) == 1
+        assert get(1, 0) == 2
+        assert get(1, 1) == 3
 
 class AppTestNumArray(BaseNumpyAppTest):
     def w_CustomIndexObject(self, index):
         s = repr(a)
         assert s.replace('\n', '') == \
                       "array(['abc', 'defg', 'ab'],       dtype='|S4')"
-         
+        
        
 class AppTestPyPy(BaseNumpyAppTest):
     def setup_class(cls):
         assert a[0][1] == 2
         a = _numpypy.array(([[[1, 2], [3, 4], [5, 6]]]))
         assert (a[0, 1] == [3, 4]).all()
+
+    def test_from_shape_and_storage(self):
+        from _numpypy import array, ndarray
+        x = array([1, 2, 3, 4])
+        addr, _ = x.__array_interface__['data']
+        y = ndarray._from_shape_and_storage([2, 2], addr, x.dtype)
+        assert y[0, 1] == 2
+        y[0, 1] = 42
+        assert x[1] == 42