[pypy-commit] pypy ufuncapi: play around with different parameters, still stuck in loop() with arbitrary function return value

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: ufuncapi
Changeset: r72140:f82d10b9c773
Date: 2014-06-22 22:37 +0300
http://bitbucket.org/pypy/pypy/changeset/f82d10b9c773/

Log:	play around with different parameters, still stuck in loop() with
	arbitrary function return value

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
@@ -127,7 +127,8 @@
         vals = [i_s[0].getitem(i_s[1]) for i_s in in_iters_and_states]
         arglist = space.wrap(vals)
         out_vals = space.call_args(func, Arguments.frompacked(space, arglist))
-        # XXX bad form
+        # XXX bad form - out_vals should be a list or tuple of boxes.
+        # but func can return anything, 
         if not isinstance(out_vals,(list, tuple)):
             out_iter, out_state = out_iters_and_states[0]
             out_iter.setitem(out_state, out_vals.convert_to(space, res_dtype))
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
@@ -201,19 +201,27 @@
         def adder(a, b):
             return a+b
         try:
-            myufunc = frompyfunc(adder, 2, 1)
+            adder_ufunc0 = frompyfunc(adder, 2, 1)
+            adder_ufunc1 = frompyfunc(adder, 2, 1)
             int_func22 = frompyfunc(int, 2, 2)
             int_func12 = frompyfunc(int, 1, 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
-            myufunc = frompyfunc(adder, 2, 1, dtypes=['match'])
-            int_func22 = frompyfunc(int, 2, 2, dtypes=['match'])
-            int_func12 = frompyfunc(int, 1, 2, dtypes=['match'])
+            adder_ufunc0 = frompyfunc(adder, 2, 1, dtypes=['match'])
+            adder_ufunc1 = frompyfunc([adder, adder], 2, 1, dtypes=[int, float])
+            int_func22 = frompyfunc([int, int], 2, 2, signature='()->()',
+                                    dtypes=[int, int, float, int])
+            int_func12 = frompyfunc([int, int], 1, 2, signature='()->()',
+                                    dtypes=[int, int, float, int])
             retype = dtype(int)
-        assert isinstance(myufunc, ufunc)
-        res = myufunc(arange(10), arange(10))
+        assert isinstance(adder_ufunc1, ufunc)
+        res = adder_ufunc0(arange(10), arange(10))
+        assert res.dtype == retype
+        assert all(res == arange(10) + arange(10))
+        res = adder_ufunc1(arange(10), arange(10))
         assert res.dtype == retype
         assert all(res == arange(10) + arange(10))
         raises(TypeError, frompyfunc, 1, 2, 3)
diff --git a/pypy/module/micronumpy/ufuncs.py b/pypy/module/micronumpy/ufuncs.py
--- a/pypy/module/micronumpy/ufuncs.py
+++ b/pypy/module/micronumpy/ufuncs.py
@@ -975,17 +975,17 @@
     else:
         raise oefmt(space.w_ValueError,
             'identity must be 0, 1, or None')
-    if nin==1 and nout==1 and dtypes == 'match':
-        w_ret = W_Ufunc1(wrap_ext_func(func[0], name))
-    elif nin==2 and nout==1 and dtypes == 'match':
-        w_ret = W_Ufunc2(wrap_ext_func(func[0], name))
+    if nin==1 and nout==1 and dtypes[0] == 'match':
+        w_ret = W_Ufunc1(wrap_ext_func(space, func[0]), name)
+    elif nin==2 and nout==1 and dtypes[0] == 'match':
+        w_ret = W_Ufunc2(wrap_ext_func(space, func[0]), name)
     else:
         w_ret = W_UfuncGeneric(space, func, name, identity, nin, nout, dtypes, signature)
     if doc:
         w_ret.w_doc = space.wrap(doc)
     return w_ret
 
-def wrap_ext_func(func):
+def wrap_ext_func(space, func):
     def _func(calc_dtype, w_left, w_right):
         arglist = space.wrap([w_left, w_right])
         return space.call_args(func, Arguments.frompacked(space, arglist))