[pypy-commit] pypy numpy-comparison: Check size of arrays in binary operations (needed for comparison)

[snus_mumrik]

Author: Ilya Osadchiy <osadchiy.ilya at gmail.com>
Branch: numpy-comparison
Changeset: r48248:949a5026cd03
Date: 2011-10-19 22:51 +0200
http://bitbucket.org/pypy/pypy/changeset/949a5026cd03/

Log:	Check size of arrays in binary operations (needed for comparison)

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
@@ -144,6 +144,7 @@
                 w_rhs.value.convert_to(calc_dtype)
             ).wrap(space)
 
+        w_lhs, w_rhs = _broadcast_arrays(space, w_lhs, w_rhs)
         new_sig = signature.Signature.find_sig([
             self.signature, w_lhs.signature, w_rhs.signature
         ])
@@ -165,6 +166,21 @@
     reduce = interp2app(W_Ufunc.descr_reduce),
 )
 
+def _broadcast_arrays(space, a1, a2):
+    from pypy.module.micronumpy.interp_numarray import Scalar
+    '''
+    Broadcast arrays to common size
+    '''
+    # For now just check sizes of two 1D arrays
+    if isinstance(a1, Scalar) or isinstance(a2, Scalar):
+        return a1, a2
+    s1 = a1.find_size()
+    s2 = a2.find_size()
+    if s1 != s2:
+        raise operationerrfmt(space.w_ValueError, "operands could not "
+            "be broadcast together with shapes (%d) (%d)", s1, s2)
+    return a1, a2
+
 def find_binop_result_dtype(space, dt1, dt2, promote_to_float=False,
     promote_bools=False):
     # dt1.num should be <= dt2.num
diff --git a/pypy/module/micronumpy/test/test_base.py b/pypy/module/micronumpy/test/test_base.py
--- a/pypy/module/micronumpy/test/test_base.py
+++ b/pypy/module/micronumpy/test/test_base.py
@@ -31,13 +31,14 @@
 
     def test_slice_signature(self, space):
         ar = SingleDimArray(10, dtype=space.fromcache(interp_dtype.W_Float64Dtype))
-        v1 = ar.descr_getitem(space, space.wrap(slice(1, 5, 1)))
+        v1 = ar.descr_getitem(space, space.wrap(slice(1, 3, 1)))
         v2 = ar.descr_getitem(space, space.wrap(slice(4, 6, 1)))
+        v3 = ar.descr_getitem(space, space.wrap(slice(3, 5, 1)))
         assert v1.signature is v2.signature
 
-        v3 = ar.descr_add(space, v1)
-        v4 = ar.descr_add(space, v2)
-        assert v3.signature is v4.signature
+        v4 = v1.descr_add(space, v2)
+        v5 = v1.descr_add(space, v3)
+        assert v4.signature is v5.signature
 
 class TestUfuncCoerscion(object):
     def test_binops(self, space):
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -604,10 +604,9 @@
         a = array(range(5))
         assert (a == None) is False
         assert (a != None) is True
-        # TODO: uncomment after size check is implemented
-        # b = array(range(2))
-        # assert (a == b) is False
-        # assert (a != b) is True
+        b = array(range(2))
+        assert (a == b) is False
+        assert (a != b) is True
 
 class AppTestSupport(object):
     def setup_class(cls):