Merge pull request #4039 from abalkin/issue-4026

BUG: ma.compressed() function returns masked array

2 files changed, 32 insertions, 4 deletions

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 68c5d5595..8dc2ca86e 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -6186,10 +6186,10 @@ def compressed(x):
         Equivalent method.
 
     """
-    if getmask(x) is nomask:
-        return np.asanyarray(x)
-    else:
-        return x.compressed()
+    if not isinstance(x, MaskedArray):
+        x = asanyarray(x)
+    return x.compressed()
+
 
 def concatenate(arrays, axis=0):
     """
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index 6608cce63..764915236 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -3308,6 +3308,34 @@ class TestMaskedArrayFunctions(TestCase):
         test = np.ma.compress(cond, marr, axis=0)
         assert_equal(test, control)
 
+    def test_compressed(self):
+        # Test ma.compressed function.
+        # Address gh-4026
+        a = np.ma.array([1, 2])
+        test = np.ma.compressed(a)
+        assert_(type(test) is np.ndarray)
+        # Test case when input data is ndarray subclass
+        class A(np.ndarray):
+            pass
+        a = np.ma.array(A(shape=0))
+        test = np.ma.compressed(a)
+        assert_(type(test) is A)
+        # Test that compress flattens
+        test = np.ma.compressed([[1],[2]])
+        assert_equal(test.ndim, 1)
+        test = np.ma.compressed([[[[[1]]]]])
+        assert_equal(test.ndim, 1)
+        # Test case when input is MaskedArray subclass
+        class M(MaskedArray):
+            pass
+        test = np.ma.compressed(M(shape=(0,1,2)))
+        assert_equal(test.ndim, 1)
+        # with .compessed() overriden
+        class M(MaskedArray):
+            def compressed(self):
+                return 42
+        test = np.ma.compressed(M(shape=(0,1,2)))
+        assert_equal(test, 42)
 
 #------------------------------------------------------------------------------
 class TestMaskedFields(TestCase):