[pypy-commit] pypy default: test, fix for broadcasting when we shouldn't

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: 
Changeset: r78978:e84ee03a93ab
Date: 2015-08-14 00:12 +0300
http://bitbucket.org/pypy/pypy/changeset/e84ee03a93ab/

Log:	test, fix for broadcasting when we shouldn't

diff --git a/pypy/module/micronumpy/concrete.py b/pypy/module/micronumpy/concrete.py
--- a/pypy/module/micronumpy/concrete.py
+++ b/pypy/module/micronumpy/concrete.py
@@ -52,11 +52,16 @@
 
     @jit.unroll_safe
     def setslice(self, space, arr):
-        if len(arr.get_shape()) > 0 and len(self.get_shape()) == 0:
-            raise oefmt(space.w_ValueError,
-                "could not broadcast input array from shape "
-                "(%s) into shape ()",
-                ','.join([str(x) for x in arr.get_shape()]))
+        if len(arr.get_shape()) >  len(self.get_shape()):
+            # record arrays get one extra dimension
+            if not self.dtype.is_record() or \
+                    len(arr.get_shape()) > len(self.get_shape()) + 1:
+                raise oefmt(space.w_ValueError,
+                    "could not broadcast input array from shape "
+                    "(%s) into shape (%s)",
+                    ','.join([str(x) for x in arr.get_shape()]),
+                    ','.join([str(x) for x in self.get_shape()]),
+                    )
         shape = shape_agreement(space, self.get_shape(), arr)
         impl = arr.implementation
         if impl.storage == self.storage:
diff --git a/pypy/module/micronumpy/test/test_ndarray.py b/pypy/module/micronumpy/test/test_ndarray.py
--- a/pypy/module/micronumpy/test/test_ndarray.py
+++ b/pypy/module/micronumpy/test/test_ndarray.py
@@ -2668,6 +2668,15 @@
         a[0, :: -1] = numpy.array([11, 12])
         assert (a == [[12, 11], [5, 10]]).all()
 
+        a = numpy.zeros((3, 2), int)
+        b = numpy.ones((3, 1), int)
+        exc = raises(ValueError, 'a[:, 1] = b')
+        assert str(exc.value) == "could not broadcast " +\
+                "input array from shape (3,1) into shape (3)"
+        a[:, 1] = b[:,0] > 0.5
+        assert (a == [[0, 1], [0, 1], [0, 1]]).all()
+        
+
     def test_ufunc(self):
         from numpy import array
         a = array([[1, 2], [3, 4], [5, 6]])