[pypy-commit] pypy numpy-refactor: sort out the agreement

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-refactor
Changeset: r57231:fc047210cc48
Date: 2012-09-07 19:32 +0200
http://bitbucket.org/pypy/pypy/changeset/fc047210cc48/

Log:	sort out the agreement

diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -65,7 +65,7 @@
         if len(args_w) > self.argcount:
             out = args_w[-1]
         else:
-            args_w = args_w[:] + [out]
+            args_w = args_w + [out]
         if out is not None and not isinstance(out, W_NDimArray):
             raise OperationError(space.w_TypeError, space.wrap(
                                             'output must be an array'))
@@ -224,13 +224,13 @@
         self.bool_result = bool_result
 
     def call(self, space, args_w):
-        if len(args_w)<2:
-            [w_obj] = args_w
+        w_obj = args_w[0]
+        if len(args_w) > 1:
+            w_out = args_w[1]
+        if space.is_w(w_out, space.w_None):
             out = None
         else:
-            [w_obj, out] = args_w
-            if space.is_w(out, space.w_None):
-                out = None
+            out = w_out
         w_obj = convert_to_array(space, w_obj)
         calc_dtype = find_unaryop_result_dtype(space,
                                   w_obj.get_dtype(),
@@ -255,7 +255,8 @@
             else:
                 out.fill(res_dtype.coerce(space, w_val))
             return out
-        shape =  shape_agreement(space, w_obj.get_shape(), out)
+        shape = shape_agreement(space, w_obj.get_shape(), out,
+                                broadcast_down=False)
         return loop.call1(shape, self.func, self.name, calc_dtype, res_dtype,
                           w_obj, out)
 
@@ -319,7 +320,7 @@
                 out = arr
             return space.wrap(out)
         new_shape = shape_agreement(space, w_lhs.get_shape(), w_rhs)
-        new_shape = shape_agreement(space, new_shape, out)
+        new_shape = shape_agreement(space, new_shape, out, broadcast_down=False)
         return loop.call2(new_shape, self.func, self.name, calc_dtype,
                           res_dtype, w_lhs, w_rhs, out)
 
diff --git a/pypy/module/micronumpy/strides.py b/pypy/module/micronumpy/strides.py
--- a/pypy/module/micronumpy/strides.py
+++ b/pypy/module/micronumpy/strides.py
@@ -104,7 +104,7 @@
             i //= shape[s]
     return coords, step, lngth
 
-def shape_agreement(space, shape1, w_arr2):
+def shape_agreement(space, shape1, w_arr2, broadcast_down=True):
     if w_arr2 is None:
         return shape1
     assert isinstance(w_arr2, W_NDimArray)
@@ -117,6 +117,13 @@
                 ",".join([str(x) for x in shape2]),
             ))
         )
+    if not broadcast_down and len([x for x in ret if x != 1]) > len([x for x in shape2 if x != 1]):
+        raise OperationError(space.w_ValueError,
+            space.wrap("unbroadcastable shape (%s) cannot be broadcasted to (%s)" % (
+                ",".join([str(x) for x in shape1]),
+                ",".join([str(x) for x in shape2]),
+            ))
+        )
     return ret
 
 def _shape_agreement(shape1, shape2):