[pypy-commit] pypy numpy-data-buffer: setting slices on buffers of arrays and slice-views

[timo_jbo]

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

Log:	setting slices on buffers of arrays and slice-views

diff --git a/pypy/module/_numpy/interp_buffer.py b/pypy/module/_numpy/interp_buffer.py
--- a/pypy/module/_numpy/interp_buffer.py
+++ b/pypy/module/_numpy/interp_buffer.py
@@ -12,18 +12,27 @@
         return self.array.get_concrete().find_size()
 
     def getitem(self, index):
+        index = self.calc_index(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[self.calc_index(index)]
+        return char_data[index]
 
     def setitem(self, index, value):
+        index = self.calc_index(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_ptr = rffi.cast(CHAR_TP, storage)
-        char_ptr[self.calc_index(index)] = value
+        char_ptr[index] = value
+
+    def setslice(self, index, newstring):
+        offset_index = self.calc_index(index)
+        if offset_index + len(newstring) > self.getlength() - 1:
+            raise IndexError("End of slice to set out of bounds (0<=index<%d)" % self.getlength())
+        for idx in range(0, len(newstring)):
+            self.setitem(index + idx, newstring[idx])
 
     def calc_index(self, index):
         return index
diff --git a/pypy/module/_numpy/test/test_buffer.py b/pypy/module/_numpy/test/test_buffer.py
--- a/pypy/module/_numpy/test/test_buffer.py
+++ b/pypy/module/_numpy/test/test_buffer.py
@@ -71,3 +71,28 @@
         assert view[0] == 4
 
         raises(IndexError, "view[4] = '\\5'")
+
+    def test_buffer_setslice(self):
+        from _numpy import array
+        from _numpy import dtype
+        ar = array(range(8), dtype=dtype("int8"))
+        buf = ar.data
+
+        buf[1:4] = '\1\1\1'
+
+        assert ar[1] == 1
+        assert ar[2] == 1
+        assert ar[3] == 1
+
+    def test_view_setslice(self):
+        from _numpy import array
+        from _numpy import dtype
+        ar = array(range(8), dtype=dtype("int8"))
+        view = ar[1:-1]
+
+        viewbuf = view.data
+        viewbuf[1:4] = '\1\1\1'
+
+        assert ar[2] == 1
+        assert ar[3] == 1
+        assert ar[4] == 1