[pypy-commit] pypy numpy-fancy-indexing: tests and fixes for bool indexing

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-fancy-indexing
Changeset: r57382:4d0afc72644d
Date: 2012-09-19 14:01 +0200
http://bitbucket.org/pypy/pypy/changeset/4d0afc72644d/

Log:	tests and fixes for bool indexing

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
@@ -64,13 +64,16 @@
 
     def getitem_filter(self, space, arr):
         if arr.get_size() > self.get_size():
-            raise OperationError(space.w_IndexError,
+            raise OperationError(space.w_ValueError,
                                  space.wrap("index out of range for array"))
         size = loop.count_all_true(arr)
         res = W_NDimArray.from_shape([size], self.get_dtype())
         return loop.getitem_filter(res, self, arr)
 
     def setitem_filter(self, space, idx, val):
+        if idx.get_size() > self.get_size():
+            raise OperationError(space.w_ValueError,
+                                 space.wrap("index out of range for array"))
         loop.setitem_filter(self, idx, val)
 
     def _prepare_array_index(self, space, w_index):
@@ -92,7 +95,7 @@
         return loop.getitem_array_int(space, self, res, iter_shape, indexes)
 
     def descr_getitem(self, space, w_idx):
-        if (isinstance(w_idx, W_NDimArray) and w_idx.get_shape() == self.get_shape() and
+        if (isinstance(w_idx, W_NDimArray) and
             w_idx.get_dtype().is_bool_type()):
             return self.getitem_filter(space, w_idx)
         try:
@@ -109,7 +112,7 @@
         self.implementation.setitem_index(space, index_list, w_value)
 
     def descr_setitem(self, space, w_idx, w_value):
-        if (isinstance(w_idx, W_NDimArray) and w_idx.get_shape() == self.get_shape() and
+        if (isinstance(w_idx, W_NDimArray) and
             w_idx.get_dtype().is_bool_type()):
             return self.setitem_filter(space, w_idx,
                                        convert_to_array(space, w_value))
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
@@ -1528,7 +1528,17 @@
         raises(IndexError, "arange(10)[[-11]]")
 
     def test_bool_array_index(self):
-        xxx
+        from numpypy import arange, array
+        b = arange(10)
+        assert (b[array([True, False, True])] == [0, 2]).all()
+        raises(ValueError, "array([1, 2])[array([True, True, True])]")
+
+    def test_bool_array_index_setitem(self):
+        from numpypy import arange, array
+        b = arange(5)
+        b[array([True, False, True])] = [20, 21]
+        assert (b == [20, 1, 21, 3, 4]).all() 
+        raises(ValueError, "array([1, 2])[array([True, False, True])] = [1, 2, 3]")
 
 class AppTestMultiDim(BaseNumpyAppTest):
     def test_init(self):