[pypy-commit] pypy setitem2d: fix for for 2d memoryview.__setitem__, (issue bb-3028)

[mattip]

Author: Matti Picus <matti.picus at gmail.com>
Branch: setitem2d
Changeset: r96829:0f9d6d72dfe9
Date: 2019-06-19 23:19 +0300
http://bitbucket.org/pypy/pypy/changeset/0f9d6d72dfe9/

Log:	fix for for 2d memoryview.__setitem__, (issue bb-3028)

diff --git a/pypy/objspace/std/memoryobject.py b/pypy/objspace/std/memoryobject.py
--- a/pypy/objspace/std/memoryobject.py
+++ b/pypy/objspace/std/memoryobject.py
@@ -132,6 +132,24 @@
         data = view.getbytes(start, itemsize)
         return view.value_from_bytes(space, data)
 
+    def _setitem_tuple_indexed(self, space, w_index, w_obj):
+        view = self.view
+        length = space.len_w(w_index)
+        ndim = view.getndim()
+        if length < ndim:
+            raise oefmt(space.w_NotImplementedError,
+                        "sub-views are not implemented")
+
+        if length > ndim:
+            raise oefmt(space.w_TypeError, \
+                    "cannot index %d-dimension view with %d-element tuple",
+                    length, ndim)
+
+        start = self._start_from_tuple(space, w_index)
+        itemsize = self.getitemsize()
+        val = self.view.bytes_from_value(space, w_obj)
+        self.view.setbytes(start * itemsize, val)
+
     def _decode_index(self, space, w_index, is_slice):
         shape = self.getshape()
         if len(shape) == 0:
@@ -184,7 +202,7 @@
         if self.view.readonly:
             raise oefmt(space.w_TypeError, "cannot modify read-only memory")
         if space.isinstance_w(w_index, space.w_tuple):
-            raise oefmt(space.w_NotImplementedError, "only 1d setitem supported")
+            return self._setitem_tuple_indexed(space, w_index, w_obj)
         start, stop, step, size = space.decode_index4(w_index, self.getlength())
         is_slice = space.isinstance_w(w_index, space.w_slice)
         start, stop, step, slicelength = self._decode_index(space, w_index, is_slice)
diff --git a/pypy/objspace/std/test/test_memoryobject.py b/pypy/objspace/std/test/test_memoryobject.py
--- a/pypy/objspace/std/test/test_memoryobject.py
+++ b/pypy/objspace/std/test/test_memoryobject.py
@@ -288,10 +288,17 @@
         assert m2.shape == m1.shape
 
     def test_2d(self):
-        m = memoryview(bytearray(b'1234123412341234')).cast('b', shape=(4, 4))
+        m = memoryview(bytearray(b'1234123412341234'))
+        assert m[3] == ord('4')
+        m[3] = ord('z')
+        assert m[3] == ord('z')
+        m = m.cast('B', shape=(4, 4))
         assert m[2, 3] == ord('4')
-        m[2, 3] = ord('z')
-        assert m[2, 3] == ord('z') 
+        m[2, 3] = ord('a')
+        assert m[2, 3] == ord('a') 
+        raises(TypeError, m.__setitem__, (2, 3), bytearray(b'12'))
+        # slices in 2d memoryviews are not supported at all
+        raises(TypeError, m.__getitem__, (slice(None), 3))
 
 class AppTestCtypes(object):
     spaceconfig = dict(usemodules=['sys', '_rawffi'])