[pypy-commit] pypy numpy-multidim: Few missing tests, code improvements

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-multidim
Changeset: r49731:5da471031a21
Date: 2011-11-24 16:48 +0200
http://bitbucket.org/pypy/pypy/changeset/5da471031a21/

Log:	Few missing tests, code improvements

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
@@ -135,6 +135,9 @@
         arr_iter = arr_iter.next(shapelen)
     return arr
 
+# Iterators for arrays
+# --------------------
+
 class BaseIterator(object):
     def next(self, shapelen):
         raise NotImplementedError
@@ -1074,10 +1077,7 @@
     def __del__(self):
         lltype.free(self.storage, flavor='raw', track_allocation=False)
 
-def zeros(space, w_size, w_dtype=None):
-    dtype = space.interp_w(interp_dtype.W_Dtype,
-        space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype)
-    )
+def _find_size_and_shape(space, w_size):
     if space.isinstance_w(w_size, space.w_int):
         size = space.int_w(w_size)
         shape = [size]
@@ -1088,22 +1088,20 @@
             item = space.int_w(w_item)
             size *= item
             shape.append(item)
+    return size, shape
+
+def zeros(space, w_size, w_dtype=None):
+    dtype = space.interp_w(interp_dtype.W_Dtype,
+        space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype)
+    )
+    size, shape = _find_size_and_shape(space, w_size)
     return space.wrap(NDimArray(size, shape[:], dtype=dtype))
 
 def ones(space, w_size, w_dtype=None):
     dtype = space.interp_w(interp_dtype.W_Dtype,
         space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype)
     )
-    if space.isinstance_w(w_size, space.w_int):
-        size = space.int_w(w_size)
-        shape = [size]
-    else:
-        size = 1
-        shape = []
-        for w_item in space.fixedview(w_size):
-            item = space.int_w(w_item)
-            size *= item
-            shape.append(item)
+    size, shape = _find_size_and_shape(space, w_size)
     arr = NDimArray(size, shape[:], dtype=dtype)
     one = dtype.adapt_val(1)
     arr.dtype.fill(arr.storage, one, 0, size)
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
@@ -908,6 +908,20 @@
         assert a.argmax() == 5
         assert a[:2,].argmax() == 3
 
+    def test_broadcast_wrong_shapes(self):
+        from numpypy import zeros
+        a = zeros((4, 3, 2))
+        b = zeros((4, 2))
+        raises(ValueError, b.__add__, a)
+
+    def test_reduce(self):
+        from numpypy import array
+        a = array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
+        assert a.sum() == (13 * 12) / 2
+        b = a[1:, 1::2]
+        c = b + b
+        assert c.sum() == (6 + 8 + 10 + 12) * 2
+
 class AppTestSupport(object):
     def setup_class(cls):
         import struct