[pypy-commit] pypy numpy-dtype-alt: (timo) ones and zeros take dtype arguments

[alex_gaynor]

Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: numpy-dtype-alt
Changeset: r46703:750c052f98bd
Date: 2011-08-22 02:40 -0500
http://bitbucket.org/pypy/pypy/changeset/750c052f98bd/

Log:	(timo) ones and zeros take dtype arguments

diff --git a/TODO.txt b/TODO.txt
--- a/TODO.txt
+++ b/TODO.txt
@@ -5,8 +5,3 @@
 * dtype guessing
 * fix sum() and prod() types
 * Any more attributes that need to be exposed at app-level
-
-For later
-=========
-
-* More APIs need dtype arguments (zeros, ones, etc.)
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
@@ -546,14 +546,20 @@
         lltype.free(self.storage, flavor='raw', track_allocation=False)
 
 @unwrap_spec(size=int)
-def zeros(space, size):
-    return space.wrap(SingleDimArray(size, dtype=space.fromcache(interp_dtype.W_Float64Dtype)))
+def zeros(space, size, w_dtype=None):
+    dtype = space.interp_w(interp_dtype.W_Dtype,
+        space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype)
+    )
+    return space.wrap(SingleDimArray(size, dtype=dtype))
 
 @unwrap_spec(size=int)
-def ones(space, size):
-    dtype = space.fromcache(interp_dtype.W_Float64Dtype)
+def ones(space, size, w_dtype=None):
+    dtype = space.interp_w(interp_dtype.W_Dtype,
+        space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype)
+    )
+
     arr = SingleDimArray(size, dtype=dtype)
-    one = dtype.box(1.0)
+    one = dtype.adapt_val(1)
     for i in xrange(size):
         arr.dtype.setitem(arr.storage, i, one)
     return space.wrap(arr)
diff --git a/pypy/module/micronumpy/test/test_dtypes.py b/pypy/module/micronumpy/test/test_dtypes.py
--- a/pypy/module/micronumpy/test/test_dtypes.py
+++ b/pypy/module/micronumpy/test/test_dtypes.py
@@ -47,3 +47,29 @@
         assert isinstance(a[0], bool)
         b = a.copy()
         assert isinstance(b[0], bool)
+
+    def test_zeros_bool(self):
+        from numpy import zeros
+        a = zeros(10, dtype=bool)
+        for i in range(10):
+            assert a[i] is False
+
+    def test_ones_bool(self):
+        from numpy import ones
+        a = ones(10, dtype=bool)
+        for i in range(10):
+            assert a[i] is True
+
+    def test_zeros_long(self):
+        from numpy import zeros
+        a = zeros(10, dtype=long)
+        for i in range(10):
+            assert isinstance(a[i], int)
+            assert a[1] == 0
+
+    def test_ones_long(self):
+        from numpy import ones
+        a = ones(10, dtype=bool)
+        for i in range(10):
+            assert isinstance(a[i], int)
+            assert a[1] == 1