[pypy-commit] pypy numpy-multidim: nonzero support and more tests for slices

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-multidim
Changeset: r48535:e7d27a90d510
Date: 2011-10-27 20:01 +0200
http://bitbucket.org/pypy/pypy/changeset/e7d27a90d510/

Log:	nonzero support and more tests for slices

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
@@ -25,7 +25,7 @@
             for w_item in space.listview(w_next):
                 stack.append(w_item)
         else:
-            w_dtype = interp_ufuncs.find_dtype_for_scalar(space, w_item, w_dtype)
+            w_dtype = interp_ufuncs.find_dtype_for_scalar(space, w_next, w_dtype)
             if w_dtype is space.fromcache(interp_dtype.W_Float64Dtype):
                 return w_dtype
     if w_dtype is None:
@@ -71,14 +71,16 @@
     def add_invalidates(self, other):
         self.invalidates.append(other)
 
-    def descr__new__(space, w_subtype, w_size_or_iterable, w_dtype=None):
+    def descr__new__(space, w_subtype, w_item_or_iterable, w_dtype=None):
         # find scalar
         if space.is_w(w_dtype, space.w_None):
-            w_dtype = _find_dtype(space, w_size_or_iterable)
+            w_dtype = _find_dtype(space, w_item_or_iterable)
         dtype = space.interp_w(interp_dtype.W_Dtype,
             space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype)
         )
-        shape, elems_w = _find_shape_and_elems(space, w_size_or_iterable)
+        if not space.issequence_w(w_item_or_iterable):
+            return scalar_w(space, dtype, w_item_or_iterable)
+        shape, elems_w = _find_shape_and_elems(space, w_item_or_iterable)
         size = len(elems_w)
         arr = NDimArray(size, shape, dtype=dtype)
         i = 0
@@ -365,6 +367,12 @@
     def descr_mean(self, space):
         return space.wrap(space.float_w(self.descr_sum(space))/self.find_size())
 
+    def descr_nonzero(self, space):
+        if self.find_size() > 1:
+            raise OperationError(space.w_ValueError, space.wrap(
+                "The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()"))
+        return self.get_concrete().eval(0).wrap(space)
+
 def convert_to_array(space, w_obj):
     if isinstance(w_obj, BaseArray):
         return w_obj
@@ -395,7 +403,10 @@
         self.value = value
 
     def find_size(self):
-        raise ValueError
+        return 1
+
+    def get_concrete(self):
+        return self
 
     def find_dtype(self):
         return self.dtype
@@ -711,6 +722,7 @@
     __pos__ = interp2app(BaseArray.descr_pos),
     __neg__ = interp2app(BaseArray.descr_neg),
     __abs__ = interp2app(BaseArray.descr_abs),
+    __nonzero__ = interp2app(BaseArray.descr_nonzero),
 
     __add__ = interp2app(BaseArray.descr_add),
     __sub__ = interp2app(BaseArray.descr_sub),
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
@@ -604,6 +604,16 @@
             for i in xrange(5):
                 assert c[i] == func(b[i], 3)
 
+    def test_nonzero(self):
+        from numpy import array
+        a = array([1, 2])
+        raises(ValueError, bool, a)
+        raises(ValueError, bool, a == a)
+        assert bool(array(1))
+        assert not bool(array(0))
+        assert bool(array([1]))
+        assert not bool(array([0]))
+
 class AppTestMultiDim(BaseNumpyAppTest):
     def test_init(self):
         import numpy
@@ -672,7 +682,7 @@
         a = numpy.array([[1, 2], [4, 5]])
         assert a[0, 1] == a[0][1] == 2
         a = numpy.array(([[[1, 2], [3, 4], [5, 6]]]))
-        assert a[0, 1] == [3, 4]
+        assert (a[0, 1] == [3, 4]).all()
 
     def test_setitem_slice(self):
         import numpy
@@ -681,11 +691,13 @@
         assert a[1, 2] == 3
         raises(TypeError, a[1].__setitem__, [1, 2, 3])
         a = numpy.array([[1, 2], [3, 4]])
-        assert a == [[1, 2], [3, 4]]
+        assert (a == [[1, 2], [3, 4]]).all()
         a[1] = numpy.array([5, 6])
-        assert a == [[1, 2], [5, 6]]
+        assert (a == [[1, 2], [5, 6]]).all()
         a[:,1] = numpy.array([8, 10])
-        assert a == [[1, 8], [5, 10]]
+        assert (a == [[1, 8], [5, 10]]).all()
+        a[:,::-1] = numpy.array([11, 12])
+        assert (a == [[12, 11], [12, 11]]).all()
 
 class AppTestSupport(object):
     def setup_class(cls):