[pypy-commit] pypy numpy-back-to-applevel: add count_reduce_numbers

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-back-to-applevel
Changeset: r51576:f53826f6f84a
Date: 2012-01-21 15:33 +0200
http://bitbucket.org/pypy/pypy/changeset/f53826f6f84a/

Log:	add count_reduce_numbers

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -28,6 +28,8 @@
         'fromstring': 'interp_support.fromstring',
         'flatiter': 'interp_numarray.W_FlatIterator',
 
+        'count_reduce_items': 'interp_numarray.count_reduce_items',
+
         'True_': 'types.Bool.True',
         'False_': 'types.Bool.False',
 
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
@@ -1308,6 +1308,23 @@
     arr.dtype.fill(arr.storage, one, 0, size)
     return space.wrap(arr)
 
+ at unwrap_spec(arr=BaseArray, skipna=bool, keepdims=bool)
+def count_reduce_items(space, arr, w_axis=None, skipna=False, keepdims=True):
+    if not keepdims:
+        raise OperationError(space.w_NotImplementedError, space.wrap("unsupported"))
+    if space.is_w(w_axis, space.w_None):
+        s = 1
+        for elem in arr.shape:
+            s *= elem
+        return space.wrap(s)
+    if space.isinstance_w(w_axis, space.w_int):
+        return space.wrap(arr.shape[space.int_w(w_axis)])
+    s = 1
+    elems = space.fixedview(w_axis)
+    for w_elem in elems:
+        s *= arr.shape[space.int_w(w_elem)]
+    return space.wrap(s)
+
 def dot(space, w_obj, w_obj2):
     w_arr = convert_to_array(space, w_obj)
     if isinstance(w_arr, Scalar):
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
@@ -1,6 +1,6 @@
 from pypy.interpreter.baseobjspace import Wrappable
 from pypy.interpreter.error import OperationError, operationerrfmt
-from pypy.interpreter.gateway import interp2app
+from pypy.interpreter.gateway import interp2app, unwrap_spec
 from pypy.interpreter.typedef import TypeDef, GetSetProperty, interp_attrproperty
 from pypy.module.micronumpy import interp_boxes, interp_dtype
 from pypy.module.micronumpy.signature import ReduceSignature,\
@@ -58,7 +58,9 @@
             )
         return self.call(space, __args__.arguments_w)
 
-    def descr_reduce(self, space, w_obj, w_dim=0):
+    @unwrap_spec(skipna=bool, keepdims=bool)
+    def descr_reduce(self, space, w_obj, w_axis=None, w_dtype=None,
+                     skipna=False, keepdims=True):
         """reduce(...)
         reduce(a, axis=0)
 
@@ -111,15 +113,18 @@
         array([[ 1,  5],
                [ 9, 13]])
         """
-        return self.reduce(space, w_obj, False, False, w_dim)
+        if space.is_w(w_axis, space.w_None):
+            axis = -1
+        else:
+            axis = space.int_w(w_axis)
+        return self.reduce(space, w_obj, False, False, axis)
 
-    def reduce(self, space, w_obj, multidim, promote_to_largest, w_dim):
+    def reduce(self, space, w_obj, multidim, promote_to_largest, dim):
         from pypy.module.micronumpy.interp_numarray import convert_to_array, \
                                                            Scalar
         if self.argcount != 2:
             raise OperationError(space.w_ValueError, space.wrap("reduce only "
                 "supported for binary functions"))
-        dim = space.int_w(w_dim)
         assert isinstance(self, W_Ufunc2)
         obj = convert_to_array(space, w_obj)
         if dim >= len(obj.shape):
@@ -494,3 +499,4 @@
 
 def get(space):
     return space.fromcache(UfuncState)
+
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -394,3 +394,11 @@
                 (3, 3.5),
             ]:
                 assert ufunc(a, b) == func(a, b)
+
+    def test_count_reduce_items(self):
+        from _numpypy import count_reduce_items, arange
+        a = arange(24).reshape(2, 3, 4)
+        assert count_reduce_items(a) == 24
+        assert count_reduce_items(a, 1) == 3
+        assert count_reduce_items(a, (1, 2)) == 3 * 4
+