[pypy-commit] pypy python-numpy: add a test, fix for numpy compatability

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: python-numpy
Changeset: r56696:b08ace68a844
Date: 2012-08-10 17:33 +0300
http://bitbucket.org/pypy/pypy/changeset/b08ace68a844/

Log:	add a test, fix for numpy compatability

diff --git a/lib_pypy/numpypy/multiarray/__init__.py b/lib_pypy/numpypy/multiarray/__init__.py
--- a/lib_pypy/numpypy/multiarray/__init__.py
+++ b/lib_pypy/numpypy/multiarray/__init__.py
@@ -41,5 +41,10 @@
 
 def count_nonzero(a):
     if not hasattr(a,'flat'):
-        a = ndarray(a)
-    return sum(a.flat != 0)
+        try:
+            a = ndarray(a)
+            return sum(a.flat != 0)
+        except TypeError:
+            if isinstance(a, (tuple, list)):
+                return len(a)
+        return 1
diff --git a/lib_pypy/numpypy/test/test_multiarray.py b/lib_pypy/numpypy/test/test_multiarray.py
new file mode 100644
--- /dev/null
+++ b/lib_pypy/numpypy/test/test_multiarray.py
@@ -0,0 +1,13 @@
+try:
+    import _numpypy as np
+    import numpypy.multiarray as multiarray
+except:
+    import numpy as np
+    from numpy.core import multiarray
+
+from py.test import raises
+
+def test_count_nonzero():
+   a = np.array([[1, 1], [1, 1]])
+   assert multiarray.count_nonzero(a) == 4
+   raises(TypeError, multiarray.count_nonzero, 'a')