[pypy-commit] pypy numpy-refactor: a test and a fix

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-refactor
Changeset: r57160:44b208ded3cb
Date: 2012-09-05 21:53 +0200
http://bitbucket.org/pypy/pypy/changeset/44b208ded3cb/

Log:	a test and a fix

diff --git a/pypy/module/micronumpy/interp_dtype.py b/pypy/module/micronumpy/interp_dtype.py
--- a/pypy/module/micronumpy/interp_dtype.py
+++ b/pypy/module/micronumpy/interp_dtype.py
@@ -18,6 +18,12 @@
 STRINGLTR = 'S'
 UNICODELTR = 'U'
 
+def decode_w_dtype(space, w_dtype):
+    if w_dtype is None or space.is_w(w_dtype, space.w_None):
+        return None
+    return space.interp_w(W_Dtype,
+          space.call_function(space.gettypefor(W_Dtype), w_dtype))
+
 class W_Dtype(Wrappable):
     _immutable_fields_ = ["itemtype", "num", "kind"]
 
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
@@ -272,7 +272,7 @@
     # ----------------------- reduce -------------------------------
 
     def _reduce_ufunc_impl(ufunc_name, promote_to_largest=False):
-        def impl(self, space, w_axis=None, w_out=None):
+        def impl(self, space, w_axis=None, w_out=None, w_dtype=None):
             if space.is_w(w_out, space.w_None) or not w_out:
                 out = None
             elif not isinstance(w_out, W_NDimArray):
@@ -282,7 +282,7 @@
                 out = w_out
             return getattr(interp_ufuncs.get(space), ufunc_name).reduce(space,
                                         self, True, promote_to_largest, w_axis,
-                                                                   False, out)
+                                                         False, out, w_dtype)
         return func_with_new_name(impl, "reduce_%s_impl" % ufunc_name)
 
     descr_sum = _reduce_ufunc_impl("add")
@@ -406,12 +406,6 @@
     tolist = interp2app(W_NDimArray.descr_tolist),
 )
 
-def decode_w_dtype(space, w_dtype):
-    if w_dtype is None or space.is_w(w_dtype, space.w_None):
-        return None
-    return space.interp_w(interp_dtype.W_Dtype,
-          space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype))
-
 @unwrap_spec(ndmin=int, copy=bool, subok=bool)
 def array(space, w_object, w_dtype=None, copy=True, w_order=None, subok=False,
           ndmin=0):
@@ -436,7 +430,7 @@
         if copy:
             return w_object.descr_copy(space)
         return w_object
-    dtype = decode_w_dtype(space, w_dtype)
+    dtype = interp_dtype.decode_w_dtype(space, w_dtype)
     shape, elems_w = find_shape_and_elems(space, w_object, dtype)
     if dtype is None:
         for w_elem in elems_w:
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
@@ -136,10 +136,11 @@
                                                 'output must be an array'))
         else:
             out = w_out
-        return self.reduce(space, w_obj, False, False, w_axis, keepdims, out)
+        return self.reduce(space, w_obj, False, False, w_axis, keepdims, out,
+                           w_dtype)
 
     def reduce(self, space, w_obj, multidim, promote_to_largest, w_axis,
-               keepdims=False, out=None):
+               keepdims=False, out=None, dtype=None):
         if self.argcount != 2:
             raise OperationError(space.w_ValueError, space.wrap("reduce only "
                 "supported for binary functions"))
@@ -153,15 +154,17 @@
         axis = unwrap_axis_arg(space, shapelen, w_axis)    
         assert axis>=0
         size = obj.get_size()
-        if self.comparison_func:
-            dtype = interp_dtype.get_dtype_cache(space).w_booldtype
-        else:
-            dtype = find_unaryop_result_dtype(
-                space, obj.get_dtype(),
-                promote_to_float=self.promote_to_float,
-                promote_to_largest=promote_to_largest,
-                promote_bools=True
-            )
+        dtype = interp_dtype.decode_w_dtype(space, dtype)
+        if dtype is None:
+            if self.comparison_func:
+                dtype = interp_dtype.get_dtype_cache(space).w_booldtype
+            else:
+                dtype = find_unaryop_result_dtype(
+                    space, obj.get_dtype(),
+                    promote_to_float=self.promote_to_float,
+                    promote_to_largest=promote_to_largest,
+                    promote_bools=True
+                )
         if self.identity is None and size == 0:
             raise operationerrfmt(space.w_ValueError, "zero-size array to "
                     "%s.reduce without identity", self.name)
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
@@ -1752,6 +1752,7 @@
         b = a[1:, 1::2]
         c = b + b
         assert c.sum() == (6 + 8 + 10 + 12) * 2
+        assert isinstance(c.sum(dtype='f8'), float)
 
     def test_transpose(self):
         from _numpypy import array