[pypy-commit] pypy numpypy_count_nonzero: Added count_nonzero to numpy

[ebsd2]

Author: Anders Lehmann <anders at hih.au.dk>
Branch: numpypy_count_nonzero
Changeset: r55988:8d399faa682a
Date: 2012-07-07 14:29 +0200
http://bitbucket.org/pypy/pypy/changeset/8d399faa682a/

Log:	Added count_nonzero to numpy

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
@@ -166,4 +166,5 @@
         'eye': 'app_numpy.eye',
         'max': 'app_numpy.max',
         'arange': 'app_numpy.arange',
+        'count_nonzero': 'app_numpy.count_nonzero',
     }
diff --git a/pypy/module/micronumpy/app_numpy.py b/pypy/module/micronumpy/app_numpy.py
--- a/pypy/module/micronumpy/app_numpy.py
+++ b/pypy/module/micronumpy/app_numpy.py
@@ -2,6 +2,10 @@
 
 import _numpypy
 
+def count_nonzero(a):
+    if not hasattr(a, 'count_nonzero'):
+        a = _numpypy.array(a)
+    return a.count_nonzero()
 
 def average(a):
     # This implements a weighted average, for now we don't implement the
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
@@ -402,6 +402,10 @@
                 i += 1
         return Chunks(result)
 
+    def descr_count_nonzero(self, space):
+        res = self.count_all_true()
+        return space.wrap(res)
+
     def count_all_true(self):
         sig = self.find_sig()
         frame = sig.create_frame(self)
@@ -1486,6 +1490,7 @@
     take = interp2app(BaseArray.descr_take),
     compress = interp2app(BaseArray.descr_compress),
     repeat = interp2app(BaseArray.descr_repeat),
+    count_nonzero = interp2app(BaseArray.descr_count_nonzero),
 )
 
 
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
@@ -2042,6 +2042,12 @@
         raises(ValueError, "array(5).item(1)")
         assert array([1]).item() == 1
 
+    def test_count_nonzero(self):
+        from _numpypy import array
+        a = array([1,0,5,0,10])
+        assert a.count_nonzero() == 3
+ 
+
 class AppTestSupport(BaseNumpyAppTest):
     def setup_class(cls):
         import struct
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
@@ -640,6 +640,13 @@
         raises(ValueError, count_reduce_items, a, -4)
         raises(ValueError, count_reduce_items, a, (0, 2, -4))
 
+    def test_count_nonzero(self):
+        from _numpypy import where, count_nonzero, arange
+        a = arange(10)
+        assert count_nonzero(a) == 9
+        a[9] = 0
+        assert count_nonzero(a) == 8
+
     def test_true_divide(self):
         from _numpypy import arange, array, true_divide
         assert (true_divide(arange(3), array([2, 2, 2])) == array([0, 0.5, 1])).all()