Merge pull request #4699 from charris/fix-ma.count-return-type-check

BUG, TST: Fix tests of ma.count return type.

2 files changed, 23 insertions, 26 deletions

diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index 02a1c8c40..e6f659041 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -808,27 +808,29 @@ class TestMaskedArrayArithmetic(TestCase):
 
     def test_count_func(self):
         # Tests count
-        ott = array([0., 1., 2., 3.], mask=[1, 0, 0, 0])
-        ott1= array([0., 1., 2., 3.])
-        if sys.version_info[0] >= 3:
-            self.assertTrue(isinstance(count(ott), np.integer))
-        else:
-            self.assertTrue(isinstance(count(ott), int))
-        assert_equal(3, count(ott))
         assert_equal(1, count(1))
         assert_equal(0, array(1, mask=[1]))
+
+        ott = array([0., 1., 2., 3.], mask=[1, 0, 0, 0])
+        res = count(ott)
+        self.assertTrue(res.dtype.type is np.intp)
+        assert_equal(3, res)
+        
         ott = ott.reshape((2, 2))
-        assert_(isinstance(count(ott, 0), ndarray))
-        if sys.version_info[0] >= 3:
-            assert_(isinstance(count(ott), np.integer))
-        else:
-            assert_(isinstance(count(ott), int))
-        assert_equal(3, count(ott))
-        assert_(getmask(count(ott, 0)) is nomask)
-        assert_equal([1, 2], count(ott, 0))
-        assert_equal(type(count(ott, 0)), type(count(ott1, 0)))
-        assert_equal(count(ott, 0).dtype, count(ott1, 0).dtype)
-        assert_raises(IndexError, ott1.count, 1)
+        res = count(ott)
+        assert_(res.dtype.type is np.intp)
+        assert_equal(3, res)
+        res = count(ott, 0)
+        assert_(isinstance(res, ndarray))
+        assert_equal([1, 2], res)
+        assert_(getmask(res) is nomask)
+
+        ott= array([0., 1., 2., 3.])
+        res = count(ott, 0)
+        assert_(isinstance(res, ndarray))
+        assert_(res.dtype.type is np.intp)
+
+        assert_raises(IndexError, ott.count, 1)
 
     def test_minmax_func(self):
         # Tests minimum and maximum.
diff --git a/numpy/ma/tests/test_old_ma.py b/numpy/ma/tests/test_old_ma.py
index 27d699385..87c2133d7 100644
--- a/numpy/ma/tests/test_old_ma.py
+++ b/numpy/ma/tests/test_old_ma.py
@@ -153,19 +153,14 @@ class TestMa(TestCase):
     def test_xtestCount(self):
         # Test count
         ott = array([0., 1., 2., 3.], mask=[1, 0, 0, 0])
-        if sys.version_info[0] >= 3:
-            self.assertTrue(isinstance(count(ott), np.integer))
-        else:
-            self.assertTrue(isinstance(count(ott), int))
+        self.assertTrue(count(ott).dtype.type is np.intp)
         self.assertEqual(3, count(ott))
         self.assertEqual(1, count(1))
         self.assertTrue(eq(0, array(1, mask=[1])))
         ott = ott.reshape((2, 2))
+        self.assertTrue(count(ott).dtype.type is np.intp)
         assert_(isinstance(count(ott, 0), np.ndarray))
-        if sys.version_info[0] >= 3:
-            assert_(isinstance(count(ott), np.integer))
-        else:
-            assert_(isinstance(count(ott), int))
+        self.assertTrue(count(ott).dtype.type is np.intp)
         self.assertTrue(eq(3, count(ott)))
         assert_(getmask(count(ott, 0)) is nomask)
         self.assertTrue(eq([1, 2], count(ott, 0)))