Merge pull request #5847 from christianbrodbeck/fix/setdiff1d

BUG: setdiff1d return type

Fixes 5846

2 files changed, 6 insertions, 6 deletions

diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index 51064e831..64a9844cf 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -1278,14 +1278,12 @@ def setdiff1d(ar1, ar2, assume_unique=False):
            fill_value = 999999)
 
     """
-    if not assume_unique:
+    if assume_unique:
+        ar1 = ma.asarray(ar1).ravel()
+    else:
         ar1 = unique(ar1)
         ar2 = unique(ar2)
-    aux = in1d(ar1, ar2, assume_unique=True)
-    if aux.size == 0:
-        return aux
-    else:
-        return ma.asarray(ar1)[aux == 0]
+    return ar1[in1d(ar1, ar2, assume_unique=True, invert=True)]
 
 
 #####--------------------------------------------------------------------------
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
index b6749ae9e..3c7b95c9e 100644
--- a/numpy/ma/tests/test_extras.py
+++ b/numpy/ma/tests/test_extras.py
@@ -1109,6 +1109,8 @@ class TestArraySetOps(TestCase):
         a = arange(10)
         b = arange(8)
         assert_equal(setdiff1d(a, b), array([8, 9]))
+        a = array([], np.uint32, mask=[])
+        assert_equal(setdiff1d(a, []).dtype, np.uint32)
 
     def test_setdiff1d_char_array(self):
         # Test setdiff1d_charray