[pypy-commit] pypy missing-ndarray-attributes: write choose tests

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: missing-ndarray-attributes
Changeset: r58583:1c008cc37599
Date: 2012-10-29 14:32 +0100
http://bitbucket.org/pypy/pypy/changeset/1c008cc37599/

Log:	write choose tests

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
@@ -401,7 +401,8 @@
             loop.byteswap(self.implementation, res.implementation)
             return res
 
-    def descr_choose(self, space, w_choices, w_out=None, w_mode='raise'):
+    @unwrap_spec(mode=str)
+    def descr_choose(self, space, w_choices, w_out=None, mode='raise'):
         raise OperationError(space.w_NotImplementedError, space.wrap(
             "choose not implemented yet"))
 
@@ -792,10 +793,11 @@
     real = GetSetProperty(W_NDimArray.descr_get_real),
     imag = GetSetProperty(W_NDimArray.descr_get_imag),
 
-    argsort = interp2app(W_NDimArray.descr_argsort),
-    astype = interp2app(W_NDimArray.descr_astype),
-    base = GetSetProperty(W_NDimArray.descr_get_base),
+    argsort  = interp2app(W_NDimArray.descr_argsort),
+    astype   = interp2app(W_NDimArray.descr_astype),
+    base     = GetSetProperty(W_NDimArray.descr_get_base),
     byteswap = interp2app(W_NDimArray.descr_byteswap),
+    choose   = interp2app(W_NDimArray.descr_choose),
 
     __array_interface__ = GetSetProperty(W_NDimArray.descr_array_iface),
 )
diff --git a/pypy/module/micronumpy/test/test_arrayops.py b/pypy/module/micronumpy/test/test_arrayops.py
--- a/pypy/module/micronumpy/test/test_arrayops.py
+++ b/pypy/module/micronumpy/test/test_arrayops.py
@@ -83,3 +83,33 @@
         assert c == 12.0
         c = array(3.0).dot(array(4))
         assert c == 12.0
+
+    def test_choose_basic(self):
+        from _numpypy import array
+        a, b, c = array([1, 2, 3]), array([4, 5, 6]), array([7, 8, 9])
+        r = array([2, 1, 0]).choose([a, b, c])
+        assert (r == [7, 5, 3]).all()
+
+    def test_choose_broadcast(self):
+        from _numpypy import array
+        a, b, c = array([1, 2, 3]), [4, 5, 6], 13
+        r = array([2, 1, 0]).choose([a, b, c])
+        assert (r == [13, 5, 3]).all()
+
+    def test_choose_out(self):
+        from _numpypy import array
+        a, b, c = array([1, 2, 3]), [4, 5, 6], 13
+        r = array([2, 1, 0]).choose([a, b, c], out=a)
+        assert (r == [13, 5, 3]).all()
+        assert (a == [13, 5, 3]).all()
+        
+    def test_choose_modes(self):
+        from _numpypy import array
+        a, b, c = array([1, 2, 3]), [4, 5, 6], 13
+        raises(ValueError, "array([3, 1, 0]).choose([a, b, c])")
+        raises(ValueError, "array([3, 1, 0]).choose([a, b, c], 'raises')")
+        r = array([4, 1, 0]).choose([a, b, c], mode='clip')
+        assert (r == [13, 5, 3]).all()
+        r = array([4, 1, 0]).choose([a, b, c], mode='wrap')
+        assert (r == [4, 5, 3]).all()
+