[pypy-commit] pypy numpy-indexing-by-arrays: bool indexing requires a numpy array and matching size. use array's dtype if it is an ind for integer indexing.

[justinpeel]

Author: Justin Peel <notmuchtotell at gmail.com>
Branch: numpy-indexing-by-arrays
Changeset: r47842:67245c6b73ee
Date: 2011-10-05 23:07 -0600
http://bitbucket.org/pypy/pypy/changeset/67245c6b73ee/

Log:	bool indexing requires a numpy array and matching size. use array's
	dtype if it is an ind for integer 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
@@ -233,25 +233,35 @@
                                      space.wrap("invalid index"))
             w_idx = space.getitem(w_idx, space.wrap(0))
         elif space.issequence_w(w_idx):
-            w_idx = convert_to_array(space, w_idx)
-            bool_dtype = space.fromcache(interp_dtype.W_BoolDtype)
-            int_dtype = space.fromcache(interp_dtype.W_Int64Dtype)
-            if w_idx.find_dtype() is bool_dtype:
-                # Indexing by boolean array
+            if isinstance(w_idx, BaseArray) and \
+                w_idx.find_dtype().kind == interp_dtype.BOOLLTR:
+                # Indexing by boolean array - must be using a bool numpy
+                # array of the exact same size as self to do this.
+                # Can't use a list.
+                bool_dtype = space.fromcache(interp_dtype.W_BoolDtype)
                 new_sig = signature.Signature.find_sig([
                     IndexedByBoolArray.signature, self.signature
-                ])                
+                ])
+                # will be more complicated with multi-dim arrays
+                if self.find_size() != w_idx.find_size():
+                    raise OperationError(space.w_ValueError,
+                        space.wrap("bool indexing requires matching dims"))
                 res = IndexedByBoolArray(new_sig, bool_dtype, self, w_idx)
                 return space.wrap(res)
             else:
+                w_idx = convert_to_array(space, w_idx)
                 # Indexing by array
-
                 # FIXME: should raise exception if any index in
                 # array is out od bound, but this kills lazy execution
                 new_sig = signature.Signature.find_sig([
                     IndexedByArray.signature, self.signature
-                ])                
-                res = IndexedByArray(new_sig, int_dtype, self, w_idx)
+                ])
+                # Use w_idx's dtype if possible
+                dtype = w_idx.find_dtype()
+                if dtype.kind != interp_dtype.SIGNEDLTR \
+                        and dtype.kind != interp_dtype.UNSIGNEDLTR:
+                    dtype = space.fromcache(interp_dtype.W_Int64Dtype)
+                res = IndexedByArray(new_sig, dtype, self, w_idx)
                 return space.wrap(res)
 
         start, stop, step, slice_length = space.decode_index4(w_idx, self.find_size())
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
@@ -136,6 +136,8 @@
     def test_index_by_bool_array(self):
         from numpy import array, dtype
         a = array(range(5))
+        ind = array([False,True])
+        raises(ValueError, "a[ind]")
         ind = array([False, True, False, True, False])
         assert ind.dtype is dtype(bool)
         # get length before actual calculation