Maciej Fijalkowski avatar Maciej Fijalkowski committed 4d0afc7

tests and fixes for bool indexing

Comments (0)

Files changed (2)

pypy/module/micronumpy/interp_numarray.py

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

pypy/module/micronumpy/test/test_numarray.py

         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):
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.