[pypy-commit] pypy numpy-singledim: numpy: added a check of array sizes to binary ops. added more tests for radd, etc.

[justinpeel]

Author: Justin Peel <notmuchtotell at gmail.com>
Branch: numpy-singledim
Changeset: r45647:dff713e0278f
Date: 2011-07-15 13:36 -0600
http://bitbucket.org/pypy/pypy/changeset/dff713e0278f/

Log:	numpy: added a check of array sizes to binary ops. added more tests
	for radd, etc.

diff --git a/pypy/module/micronumpy/compile.py b/pypy/module/micronumpy/compile.py
--- a/pypy/module/micronumpy/compile.py
+++ b/pypy/module/micronumpy/compile.py
@@ -4,6 +4,7 @@
 """
 
 from pypy.module.micronumpy.interp_numarray import FloatWrapper, SingleDimArray
+from pypy.rlib.objectmodel import specialize
 
 class BogusBytecode(Exception):
     pass
@@ -15,8 +16,12 @@
     return a
 
 class TrivialSpace(object):
-    def wrap(self, x):
-        return x
+    w_ValueError = None
+
+    @specialize.argtype(1)
+    def wrap(self, w_obj):
+        return w_obj
+
 
 def numpy_compile(bytecode, array_size):
     space = TrivialSpace()
diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -90,6 +90,18 @@
         signature = Signature()
         def impl(self, space, w_other):
             w_other = convert_to_array(space, w_other)
+            try:
+                w_other_size = w_other.find_size()
+                self_size = self.find_size()
+            except ValueError:
+                # this will be raised if one of the arrays is a scalar.
+                pass
+            else:
+                # Need a better dimension check here for N-dim arrays
+                if w_other_size != self_size:
+                    raise OperationError(space.w_ValueError,
+                        space.wrap("Cannot %s arrays of unequal dimensions" \
+                        % function.__name__))
             new_sig = self.signature.transition(signature)
             res = Call2(
                 function,
@@ -113,7 +125,7 @@
         signature = Signature()
         def impl(self, space, w_other):
             new_sig = self.signature.transition(signature)
-            w_other = FloatWrapper(space.float_w(w_other))
+            w_other = convert_to_array(space, w_other)
             res = Call2(
                 function,
                 w_other,
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
@@ -18,8 +18,8 @@
 
     def test_slice_signature(self, space):
         ar = SingleDimArray(10)
-        v1 = ar.descr_getitem(space, space.wrap(slice(1, 5, 1)))
-        v2 = ar.descr_getitem(space, space.wrap(slice(4, 6, 1)))
+        v1 = ar.descr_getitem(space, space.wrap(slice(0, 10, 1)))
+        v2 = ar.descr_getitem(space, space.wrap(slice(9, None, -1)))
         assert v1.signature is v2.signature
 
         v3 = ar.descr_add(space, v1)
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
@@ -145,6 +145,9 @@
         b = a + 5
         for i in range(5):
             assert b[i] == i + 5
+        b = 5 + a
+        for i in range(5):
+            assert b[i] == 5 + i
 
     def test_add_list(self):
         from numpy import array
@@ -154,6 +157,16 @@
         assert isinstance(c, array)
         for i in range(5):
             assert c[i] == 4
+        c = b + a
+        assert isinstance(c, array)
+        for i in range(5):
+            assert c[i] == 4
+
+    def test_add_unequal_size(self):
+        from numpy import array
+        a = array(range(5))
+        b = array(range(3))
+        raises(ValueError, "a + b")
 
     def test_subtract(self):
         from numpy import array
@@ -176,6 +189,9 @@
         b = a - 5
         for i in range(5):
             assert b[i] == i - 5
+        b = 5 - a
+        for i in range(5):
+            assert b[i] == 5 - i
 
     def test_mul(self):
         from numpy import array
@@ -190,6 +206,9 @@
         b = a * 5
         for i in range(5):
             assert b[i] == i * 5
+        b = 5 * a
+        for i in range(5):
+            assert b[i] == 5 * i
 
     def test_div(self):
         from numpy import array
@@ -208,10 +227,13 @@
 
     def test_div_constant(self):
         from numpy import array
-        a = array(range(5))
+        a = array(range(1,6))
         b = a / 5.0
         for i in range(5):
-            assert b[i] == i / 5.0
+            assert b[i] == (i+1) / 5.0
+        b = 5.0 / a
+        for i in range(5):
+            assert b[i] == 5.0 / (i+1)
 
     def test_pow(self):
         from numpy import array
@@ -235,6 +257,9 @@
         b = a ** 2
         for i in range(5):
             assert b[i] == i ** 2
+        b = 2 ** a
+        for i in range(5):
+            assert b[i] == 2 ** i
 
     def test_mod(self):
         from numpy import array
@@ -253,10 +278,13 @@
 
     def test_mod_constant(self):
         from numpy import array
-        a = array(range(5))
+        a = array(range(1,6))
         b = a % 2
         for i in range(5):
-            assert b[i] == i % 2
+            assert b[i] == (i+1) % 2
+        b = 2 % a
+        for i in range(5):
+            assert b[i] == 2 % (i+1)
 
     def test_pos(self):
         from numpy import array