[pypy-commit] pypy default: test, fix indexing with a list of slices

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: 
Changeset: r62317:6702e78d079d
Date: 2013-03-11 17:37 -0500
http://bitbucket.org/pypy/pypy/changeset/6702e78d079d/

Log:	test, fix indexing with a list of slices

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
@@ -136,6 +136,10 @@
     def getitem_array_int(self, space, w_index):
         prefix, res_shape, iter_shape, indexes = \
                 self._prepare_array_index(space, w_index)
+        if iter_shape is None:
+            # w_index is a list of slices, return a view
+            chunks = self.implementation._prepare_slice_args(space, w_index)
+            return chunks.apply(self)
         shape = res_shape + self.get_shape()[len(indexes):]
         res = W_NDimArray.from_shape(shape, self.get_dtype(), self.get_order())
         if not res.get_size():
@@ -147,6 +151,13 @@
         val_arr = convert_to_array(space, w_value)
         prefix, _, iter_shape, indexes = \
                 self._prepare_array_index(space, w_index)
+        if iter_shape is None:
+            # w_index is a list of slices
+            w_value = convert_to_array(space, w_value)
+            chunks = self.implementation._prepare_slice_args(space, w_index)
+            view = chunks.apply(self)
+            view.implementation.setslice(space, w_value)
+            return
         return loop.setitem_array_int(space, self, iter_shape, indexes, val_arr,
                                       prefix)
 
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
@@ -2204,6 +2204,20 @@
         a[array([0, 2]), slice(0, 2)] = [[10, 11], [12, 13]]
         assert (a == [[10, 11], [3, 4], [12, 13]]).all()
 
+    def test_slice_vector_index(self):
+        from numpypy import arange
+        b = arange(145)
+        a = b[slice(25, 125, None)]
+        assert (a == range(25, 125)).all()
+        a = b[[slice(25, 125, None)]]
+        assert a.shape == (100,)
+        # a is a view into b
+        a[10] = 200
+        assert b[35] == 200
+        b[[slice(25, 30)]] = range(5)
+        assert all(a[:5] == range(5))
+        raises(TypeError, 'b[[[slice(25, 125)]]]')
+
     def test_cumsum(self):
         from numpypy import arange
         a = arange(6).reshape(3, 2)