[pypy-commit] pypy separate-applevel-numpy: implement creating arrays with tuples as size arguments.

[timo_jbo]

Author: Timo Paulssen <timonator at perpetuum-immobile.de>
Branch: separate-applevel-numpy
Changeset: r47639:a38ccdc30d81
Date: 2011-09-27 22:03 +0200
http://bitbucket.org/pypy/pypy/changeset/a38ccdc30d81/

Log:	implement creating arrays with tuples as size arguments.

diff --git a/pypy/module/_numpy/interp_numarray.py b/pypy/module/_numpy/interp_numarray.py
--- a/pypy/module/_numpy/interp_numarray.py
+++ b/pypy/module/_numpy/interp_numarray.py
@@ -550,19 +550,32 @@
     def __del__(self):
         lltype.free(self.storage, flavor='raw', track_allocation=False)
 
- at unwrap_spec(size=int)
-def zeros(space, size, w_dtype=None):
+def size_w_to_single_dim_size(space, w_size):
+    if space.isinstance_w(w_size, space.w_tuple):
+        if space.len_w(w_size) != 1:
+            raise OperationError(space.w_ValueError,
+                                 space.wrap("Invalid number of dimensions"))
+        w_size_int = space.getitem(w_size, space.wrap(0))
+        size = space.unwrap(w_size_int)
+    else:
+        size = space.unwrap(w_size)
+
+    return size
+
+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 = size_w_to_single_dim_size(space, w_size)
     return space.wrap(SingleDimArray(size, dtype=dtype))
 
- at unwrap_spec(size=int)
-def ones(space, size, w_dtype=None):
+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)
     )
 
+    size = size_w_to_single_dim_size(space, w_size)
+
     arr = SingleDimArray(size, dtype=dtype)
     one = dtype.adapt_val(1)
     for i in xrange(size):
diff --git a/pypy/module/_numpy/test/test_numarray.py b/pypy/module/_numpy/test/test_numarray.py
--- a/pypy/module/_numpy/test/test_numarray.py
+++ b/pypy/module/_numpy/test/test_numarray.py
@@ -17,6 +17,15 @@
         a[13] = 5.3
         assert a[13] == 5.3
 
+    def test_zeros_tuple_shape(self):
+        from numpy import zeros, ones, empty
+        a = zeros((15,))
+        assert len(a) == 15
+        b = ones((3,))
+        assert len(b) == 3
+        c = empty((2,))
+        assert len(c) == 2
+
     def test_empty(self):
         """
         Test that empty() works.