Fix uses of nonzero and add flatnonzero

3 files changed, 15 insertions, 3 deletions

diff --git a/numpy/lib/convertnumericA.py b/numpy/lib/convertnumericA.py
index ad1daa329..c3748bae3 100644
--- a/numpy/lib/convertnumericA.py
+++ b/numpy/lib/convertnumericA.py
@@ -13,6 +13,9 @@ Makes the following changes:
  * Converts .flat to .ravel() except for .flat = xxx or .flat[xxx]
  * Replace xxx.spacesaver() with True
  * Convert xx.savespace(?) to pass + ## xx.savespace(?)
+
+ * Converts uses of 'b' to 'B' in the typecode-position of
+   functions and methods
 """
 __all__ = ['fromfile', 'fromstr']
 
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 557383b71..548af953d 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -7,7 +7,8 @@ __all__ = ['logspace', 'linspace',
            'nanargmin', 'nanmin', 'vectorize', 'asarray_chkfinite', 'average',
            'histogram', 'bincount', 'digitize', 'cov', 'corrcoef', 'msort',
            'median', 'sinc', 'hamming', 'hanning', 'bartlett', 'blackman',
-           'kaiser', 'trapz', 'i0', 'add_newdoc', 'add_docstring', 'meshgrid'
+           'kaiser', 'trapz', 'i0', 'add_newdoc', 'add_docstring', 'meshgrid',
+           'flatnonzero'
            ]
 
 import types
@@ -530,7 +531,7 @@ def extract(condition, arr):
 
     Equivalent to compress(ravel(condition), ravel(arr)).
     """
-    return _nx.take(ravel(arr), nonzero(ravel(condition)))
+    return _nx.take(ravel(arr), nonzero(ravel(condition))[0])
 
 def insert(arr, mask, vals):
     """Similar to putmask arr[mask] = vals but the 1D array vals has the
@@ -997,3 +998,11 @@ def meshgrid(x,y):
     y = y.reshape(numRows,1)
     Y = y.repeat(numCols, axis=1)
     return X, Y
+
+def flatnonzero(a):
+    """Return indicies that are not-zero in flattened version of a
+
+    Equivalent to a.ravel().nonzero()[0]
+    """
+    return a.ravel().nonzero()[0]
+    
diff --git a/numpy/lib/polynomial.py b/numpy/lib/polynomial.py
index 206558179..0d66fd6b0 100644
--- a/numpy/lib/polynomial.py
+++ b/numpy/lib/polynomial.py
@@ -90,7 +90,7 @@ def roots(p):
         raise ValueError,"Input must be a rank-1 array."
 
     # find non-zero array entries
-    non_zero = NX.nonzero(NX.ravel(p))
+    non_zero = NX.nonzero(NX.ravel(p))[0]
 
     # find the number of trailing zeros -- this is the number of roots at 0.
     trailing_zeros = len(p) - non_zero[-1] - 1