Add 'compress'.

2 files changed, 14 insertions, 2 deletions

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index bb7e339fc..2ce1a2bd5 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -26,8 +26,8 @@ __all__ = ['MAError', 'MaskType', 'MaskedArray',
            'arctanh', 'argmax', 'argmin', 'argsort', 'around',
            'array', 'asarray','asanyarray',
            'bitwise_and', 'bitwise_or', 'bitwise_xor',
-           'ceil', 'choose', 'compressed', 'concatenate', 'conjugate',
-           'cos', 'cosh', 'count',
+           'ceil', 'choose', 'compress', 'compressed', 'concatenate',
+           'conjugate', 'cos', 'cosh', 'count',
            'default_fill_value', 'diagonal', 'divide', 'dump', 'dumps',
            'empty', 'empty_like', 'equal', 'exp',
            'fabs', 'fmod', 'filled', 'floor', 'floor_divide','fix_invalid',
@@ -3099,6 +3099,12 @@ def choose (indices, t, out=None, mode='raise'):
     m = make_mask(mask_or(m, getmask(indices)), copy=0, shrink=True)
     return masked_array(d, mask=m)
 
+def compress(a, condition):
+    """Return a where condition is True.
+
+    """
+    return a[condition]
+
 def round_(a, decimals=0, out=None):
     """Return a copy of a, rounded to 'decimals' places.
 
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index 2dc758513..9b4f19284 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -1388,6 +1388,12 @@ class TestArrayMethods(NumpyTestCase):
         putmask(mxx, mask, values)
         assert_equal(mxx, [1,2,30,4,5,60])
 
+    def test_compress(self):
+        a = array([1,2,3],mask=[True,False,False])
+        b = compress(a,a<3)
+        assert_equal(b,[1,2])
+        assert_equal(b.mask,[True,False])
+
 
 #..............................................................................