[pypy-commit] pypy numpy-data-buffer: first bit of buffer support, supports reading, somewhat.

[timo_jbo]

Author: Timo Paulssen <timonator at perpetuum-immobile.de>
Branch: numpy-data-buffer
Changeset: r47774:103fcb594b05
Date: 2011-10-03 05:12 +0200
http://bitbucket.org/pypy/pypy/changeset/103fcb594b05/

Log:	first bit of buffer support, supports reading, somewhat.

diff --git a/pypy/module/_numpy/interp_buffer.py b/pypy/module/_numpy/interp_buffer.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/_numpy/interp_buffer.py
@@ -0,0 +1,20 @@
+from pypy.interpreter.buffer import RWBuffer
+from pypy.rpython.lltypesystem import lltype, rffi
+
+CHAR_TP = lltype.Ptr(lltype.Array(lltype.Char, hints={'nolength': True}))
+
+class NumpyBuffer(RWBuffer):
+    def __init__(self, array):
+        RWBuffer.__init__(self)
+        self.array = array
+
+    def getlength(self):
+        return self.array.get_concrete().find_size()
+
+    def getitem(self, index):
+        if index > self.getlength() - 1:
+            raise IndexError("Index out of bounds (0<=index<%d)" % self.getlength())
+        storage = self.array.get_concrete().get_root_storage()
+        char_data = rffi.cast(CHAR_TP, storage)
+        return char_data[index]
+
diff --git a/pypy/module/_numpy/interp_numarray.py b/pypy/module/_numpy/interp_numarray.py
--- a/pypy/module/_numpy/interp_numarray.py
+++ b/pypy/module/_numpy/interp_numarray.py
@@ -7,6 +7,7 @@
 from pypy.rpython.lltypesystem import lltype
 from pypy.tool.sourcetools import func_with_new_name
 
+from pypy.module._numpy.interp_buffer import NumpyBuffer
 
 numpy_driver = jit.JitDriver(greens = ['signature'],
                              reds = ['result_size', 'i', 'self', 'result'])
@@ -19,6 +20,7 @@
 
     def __init__(self):
         self.invalidates = []
+        self._buffer = None
 
     def invalidated(self):
         if self.invalidates:
@@ -292,6 +294,11 @@
             j += 1
             i += step
 
+    def descr_get_data(self, space):
+        if self._buffer is None:
+            self._buffer = NumpyBuffer(self)
+        return space.wrap(self._buffer)
+
 def convert_to_array(space, w_obj):
     if isinstance(w_obj, BaseArray):
         return w_obj
@@ -622,6 +629,8 @@
     dtype = GetSetProperty(BaseArray.descr_get_dtype),
     shape = GetSetProperty(BaseArray.descr_get_shape),
 
+    data = GetSetProperty(BaseArray.descr_get_data),
+
     mean = interp2app(BaseArray.descr_mean),
     sum = interp2app(BaseArray.descr_sum),
     prod = interp2app(BaseArray.descr_prod),
diff --git a/pypy/module/_numpy/test/test_buffer.py b/pypy/module/_numpy/test/test_buffer.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/_numpy/test/test_buffer.py
@@ -0,0 +1,41 @@
+from pypy.module._numpy.test.test_base import BaseNumpyAppTest
+
+class AppTestNumArray(BaseNumpyAppTest):
+    def test_access(self):
+        from _numpy import array
+        from _numpy import dtype
+        ar = array(range(5), dtype=dtype("int8"))
+        buf = ar.data
+
+        assert buf[0] == '\0'
+        assert buf[1] == '\1'
+
+        raises(IndexError, "buf[5]")
+
+    def test_mutable(self):
+        from _numpy import array
+        from _numpy import dtype
+        ar = array(range(5), dtype=dtype("int8"))
+        buf = ar.data
+        assert buf[0] == '\0'
+
+        ar[0] = 5
+        assert buf[0] == "\5"
+
+    def test_slice_view(self):
+        skip("buffers on slicing views doesn't work yet")
+        from _numpy import array
+        from _numpy import dtype
+        ar = array(range(5), dtype=dtype("int8"))
+
+        view = ar[1:-1]
+
+        arbuf = ar.data
+        viewbuf = view.data
+
+        ar[1] = 5
+
+        assert ar[1] == view[0] == 5
+
+        assert arbuf[1] == '\5'
+        assert viewbuf[0] == '\5'