[pypy-commit] pypy ufuncapi: test, fix non-rpython many-output-ufunc, (fijal reviewing)

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: ufuncapi
Changeset: r75201:b4d4ba49c7bb
Date: 2015-01-01 23:38 +0200
http://bitbucket.org/pypy/pypy/changeset/b4d4ba49c7bb/

Log:	test, fix non-rpython many-output-ufunc, (fijal reviewing)

diff --git a/pypy/module/micronumpy/loop.py b/pypy/module/micronumpy/loop.py
--- a/pypy/module/micronumpy/loop.py
+++ b/pypy/module/micronumpy/loop.py
@@ -126,9 +126,9 @@
                                      res_dtype=res_dtype, nin=nin)
         for i in range(nin):
             vals[i] = in_iters[i].getitem(in_states[i])
-        arglist = space.newlist(vals)
-        out_val = space.call_args(func, Arguments.frompacked(space, arglist))
-        out_iter.setitem(out_state, res_dtype.coerce(space, out_val))
+        w_arglist = space.newlist(vals)
+        w_out_val = space.call_args(func, Arguments.frompacked(space, w_arglist))
+        out_iter.setitem(out_state, res_dtype.coerce(space, w_out_val))
         for i in range(nin):
             in_states[i] = in_iters[i].next(in_states[i])
         out_state = out_iter.next(out_state)
@@ -162,23 +162,22 @@
         out_states[i] = out_state
     shapelen = len(shape)
     vals = [None] * nin
-    # what does the function return?
     while not out_iters[0].done(out_states[0]):
         call_many_to_many_driver.jit_merge_point(shapelen=shapelen, func=func,
                                      res_dtype=res_dtype, nin=nin, nout=nout)
         for i in range(nin):
             vals[i] = in_iters[i].getitem(in_states[i])
-        arglist = space.newlist(vals)
-        out_vals = space.call_args(func, Arguments.frompacked(space, arglist))
-        # XXX bad form - out_vals should be a list or tuple of boxes.
-        # but func can return anything, 
-        if not isinstance(out_vals, list) and not isinstance(out_vals, tuple):
-            out_iters[0].setitem(out_states[0], res_dtype.coerce(space, out_vals))
+        w_arglist = space.newlist(vals)
+        w_outvals = space.call_args(func, Arguments.frompacked(space, w_arglist))
+        # w_outvals should be a tuple, but func can return a single value as well 
+        if space.isinstance_w(w_outvals, space.w_tuple):
+            batch = space.listview(w_outvals)
+            for i in range(len(batch)):
+                out_iters[i].setitem(out_states[i], res_dtype.coerce(space, batch[i]))
+                out_states[i] = out_iters[i].next(out_states[i])
+        else:
+            out_iters[0].setitem(out_states[0], res_dtype.coerce(space, w_outvals))
             out_states[0] = out_iters[0].next(out_states[0])
-        else:    
-            for i in range(len(out_vals)):
-                out_iters[i].setitem(out_states[i], res_dtype.coerce(space, out_vals[i]))
-                out_states[i] = out_iters[i].next(out_states[i])
         for i in range(nin):
             in_states[i] = in_iters[i].next(in_states[i])
     return space.newtuple([convert_to_array(space, o) for o in out_args])
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -149,17 +149,22 @@
         from numpy import ufunc, frompyfunc, arange, dtype
         def adder(a, b):
             return a+b
+        def sumdiff(a, b):
+                return a+b, a-b
         try:
             adder_ufunc0 = frompyfunc(adder, 2, 1)
             adder_ufunc1 = frompyfunc(adder, 2, 1)
             int_func22 = frompyfunc(int, 2, 2)
             int_func12 = frompyfunc(int, 1, 2)
+            sumdiff = frompyfunc(sumdiff, 2, 2)
             retype = dtype(object)
         except NotImplementedError as e:
             # dtype of returned value is object, which is not supported yet
             assert 'object' in str(e)
             # Use pypy specific extension for out_dtype
             adder_ufunc0 = frompyfunc(adder, 2, 1, dtypes=['match'])
+            sumdiff = frompyfunc(sumdiff, 2, 2, dtypes=['match'], 
+                                    signature='(i),(i)->(i),(i)')
             adder_ufunc1 = frompyfunc([adder, adder], 2, 1,
                             dtypes=[int, int, int, float, float, float])
             int_func22 = frompyfunc([int, int], 2, 2, signature='(i),(i)->(i),(i)',
@@ -167,19 +172,23 @@
             int_func12 = frompyfunc([int], 1, 2, signature='(i)->(i),(i)',
                                     dtypes=['match'])
             retype = dtype(int)
+        a = arange(10)
         assert isinstance(adder_ufunc1, ufunc)
-        res = adder_ufunc0(arange(10), arange(10))
+        res = adder_ufunc0(a, a)
         assert res.dtype == retype
-        assert all(res == arange(10) + arange(10))
-        res = adder_ufunc1(arange(10), arange(10))
+        assert all(res == a + a)
+        res = adder_ufunc1(a, a)
         assert res.dtype == retype
-        assert all(res == arange(10) + arange(10))
+        assert all(res == a + a)
         raises(TypeError, frompyfunc, 1, 2, 3)
-        raises (ValueError, int_func22, arange(10))
-        res = int_func12(arange(10))
+        raises (ValueError, int_func22, a)
+        res = int_func12(a)
         assert len(res) == 2
         assert isinstance(res, tuple)
-        assert (res[0] == arange(10)).all()
+        assert (res[0] == a).all()
+        res = sumdiff(2 * a, a)
+        assert (res[0] == 3 * a).all()
+        assert (res[1] == a).all()
 
     def test_frompyfunc_outerloop(self):
         def int_times2(in_array, out_array):