[pypy-commit] pypy default: Fix more of fancy indexing

[rguillebert]

Author: Romain Guillebert <romain.py at gmail.com>
Branch: 
Changeset: r66378:dbb60d9a195a
Date: 2013-08-28 13:36 +0100
http://bitbucket.org/pypy/pypy/changeset/dbb60d9a195a/

Log:	Fix more of fancy 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
@@ -88,7 +88,9 @@
         w_res = W_NDimArray.from_shape(space, res_shape, self.get_dtype(), w_instance=self)
         return loop.getitem_filter(w_res, self, arr)
 
-    def setitem_filter(self, space, idx, val):
+    def setitem_filter(self, space, idx, value):
+        from pypy.module.micronumpy.interp_boxes import Box
+        val = value
         if len(idx.get_shape()) > 1 and idx.get_shape() != self.get_shape():
             raise OperationError(space.w_ValueError,
                                  space.wrap("boolean index array should have 1 dimension"))
@@ -97,10 +99,16 @@
                                  space.wrap("index out of range for array"))
         idx_iter = idx.create_iter(self.get_shape())
         size = loop.count_all_true_iter(idx_iter, self.get_shape(), idx.get_dtype())
-        if size != val.get_shape()[0]:
+        if len(val.get_shape()) > 0 and val.get_shape()[0] > 1 and size > val.get_shape()[0]:
             raise OperationError(space.w_ValueError, space.wrap("NumPy boolean array indexing assignment "
                                                                 "cannot assign %d input values to "
                                                                 "the %d output values where the mask is true" % (val.get_shape()[0],size)))
+        if val.get_shape() == [1]:
+            box = val.descr_getitem(space, space.wrap(0))
+            assert isinstance(box, Box)
+            val = W_NDimArray(scalar.Scalar(val.get_dtype(), box))
+        elif val.get_shape() == [0]:
+            val.implementation.dtype = self.implementation.dtype
         loop.setitem_filter(self, idx, val)
 
     def _prepare_array_index(self, space, w_index):