[pypy-commit] pypy fix-result-types: simplify loop.call1()

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: fix-result-types
Changeset: r77335:3319feb32a6a
Date: 2015-05-14 17:07 +0100
http://bitbucket.org/pypy/pypy/changeset/3319feb32a6a/

Log:	simplify loop.call1()

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
@@ -91,16 +91,12 @@
     greens=['shapelen', 'func', 'calc_dtype', 'res_dtype'],
     reds='auto')
 
-def call1(space, shape, func, calc_dtype, res_dtype, w_obj, out):
+def call1(space, shape, func, calc_dtype, w_obj, w_ret):
     obj_iter, obj_state = w_obj.create_iter(shape)
     obj_iter.track_index = False
-
-    if out is None:
-        w_ret = W_NDimArray.from_shape(space, shape, res_dtype, w_instance=w_obj)
-    else:
-        w_ret = out
     out_iter, out_state = w_ret.create_iter(shape)
     shapelen = len(shape)
+    res_dtype = w_ret.get_dtype()
     while not out_iter.done(out_state):
         call1_driver.jit_merge_point(shapelen=shapelen, func=func,
                                      calc_dtype=calc_dtype, res_dtype=res_dtype)
@@ -108,8 +104,6 @@
         out_iter.setitem(out_state, func(calc_dtype, elem).convert_to(space, res_dtype))
         out_state = out_iter.next(out_state)
         obj_state = obj_iter.next(obj_state)
-    if out is None:
-        w_ret = space.call_method(w_obj, '__array_wrap__', w_ret)
     return w_ret
 
 call_many_to_one_driver = jit.JitDriver(
@@ -181,7 +175,7 @@
             vals[i] = in_iters[i].getitem(in_states[i])
         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 
+        # 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)):
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
@@ -413,8 +413,15 @@
         assert isinstance(w_obj, W_NDimArray)
         shape = shape_agreement(space, w_obj.get_shape(), out,
                                 broadcast_down=False)
-        return loop.call1(space, shape, func, calc_dtype, res_dtype,
-                          w_obj, out)
+        if out is None:
+            w_res = W_NDimArray.from_shape(
+                space, shape, res_dtype, w_instance=w_obj)
+        else:
+            w_res = out
+        w_res = loop.call1(space, shape, func, calc_dtype, w_obj, w_res)
+        if out is None:
+            w_res = space.call_method(w_obj, '__array_wrap__', w_res)
+        return w_res
 
     def call_scalar(self, space, w_arg, in_dtype, out_dtype, out):
         w_val = self.func(in_dtype, w_arg.convert_to(space, in_dtype))