Commits

mattip  committed 6702e78

test, fix indexing with a list of slices

  • Participants
  • Parent commits 241e17f

Comments (0)

Files changed (2)

File pypy/module/micronumpy/interp_numarray.py

     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():
         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)
 

File pypy/module/micronumpy/test/test_numarray.py

         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)