TST: Expand/rewrite the `test_allnans` tests

Parametrize w.r.t. the axis, array dimensionality and dtype

1 files changed, 115 insertions, 79 deletions

diff --git a/numpy/lib/tests/test_nanfunctions.py b/numpy/lib/tests/test_nanfunctions.py
index fd5e9599c..931b32211 100644
--- a/numpy/lib/tests/test_nanfunctions.py
+++ b/numpy/lib/tests/test_nanfunctions.py
@@ -83,21 +83,23 @@ class TestNanFunctions_MinMax:
             res = nf(_ndat, axis=1)
             assert_almost_equal(res, tgt)
 
-    def test_allnans(self):
-        mat = np.array([np.nan]*9).reshape(3, 3)
-        for f in self.nanfuncs:
-            for axis in [None, 0, 1]:
-                with warnings.catch_warnings(record=True) as w:
-                    warnings.simplefilter('always')
-                    assert_(np.isnan(f(mat, axis=axis)).all())
-                    assert_(len(w) == 1, 'no warning raised')
-                    assert_(issubclass(w[0].category, RuntimeWarning))
-            # Check scalars
-            with warnings.catch_warnings(record=True) as w:
-                warnings.simplefilter('always')
-                assert_(np.isnan(f(np.nan)))
-                assert_(len(w) == 1, 'no warning raised')
-                assert_(issubclass(w[0].category, RuntimeWarning))
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"])
+    @pytest.mark.parametrize("array", [
+        np.full((3, 3), np.nan),
+        np.array(np.nan),
+    ], ids=["0d", "2d"])
+    def test_allnans(self, axis, dtype, array):
+        if axis is not None and array.ndim == 0:
+            pytest.skip(f"`axis != None` not supported for 0d arrays")
+
+        array = array.astype(dtype)
+        match = "All-NaN slice encountered"
+        for func in self.nanfuncs:
+            with pytest.warns(RuntimeWarning, match=match):
+                out = func(array, axis=axis)
+            assert np.isnan(out).all()
+            assert out.dtype == array.dtype
 
     def test_masked(self):
         mat = np.ma.fix_invalid(_ndat)
@@ -193,12 +195,20 @@ class TestNanFunctions_ArgminArgmax:
                     assert_(not fcmp(val, row).any())
                     assert_(not np.equal(val, row[:ind]).any())
 
-    def test_allnans(self):
-        mat = np.array([np.nan]*9).reshape(3, 3)
-        for f in self.nanfuncs:
-            for axis in [None, 0, 1]:
-                assert_raises(ValueError, f, mat, axis=axis)
-            assert_raises(ValueError, f, np.nan)
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"])
+    @pytest.mark.parametrize("array", [
+        np.full((3, 3), np.nan),
+        np.array(np.nan),
+    ], ids=["0d", "2d"])
+    def test_allnans(self, axis, dtype, array):
+        if axis is not None and array.ndim == 0:
+            pytest.skip(f"`axis != None` not supported for 0d arrays")
+
+        array = array.astype(dtype)
+        for func in self.nanfuncs:
+            with pytest.raises(ValueError, match="All-NaN slice encountered"):
+                func(array, axis=axis)
 
     def test_empty(self):
         mat = np.zeros((0, 3))
@@ -417,20 +427,21 @@ class TestNanFunctions_SumProd(SharedNanFunctionsTestsMixin):
     nanfuncs = [np.nansum, np.nanprod]
     stdfuncs = [np.sum, np.prod]
 
-    def test_allnans(self):
-        # Check for FutureWarning
-        with warnings.catch_warnings(record=True) as w:
-            warnings.simplefilter('always')
-            res = np.nansum([np.nan]*3, axis=None)
-            assert_(res == 0, 'result is not 0')
-            assert_(len(w) == 0, 'warning raised')
-            # Check scalar
-            res = np.nansum(np.nan)
-            assert_(res == 0, 'result is not 0')
-            assert_(len(w) == 0, 'warning raised')
-            # Check there is no warning for not all-nan
-            np.nansum([0]*3, axis=None)
-            assert_(len(w) == 0, 'unwanted warning raised')
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"])
+    @pytest.mark.parametrize("array", [
+        np.full((3, 3), np.nan),
+        np.array(np.nan),
+    ], ids=["0d", "2d"])
+    def test_allnans(self, axis, dtype, array):
+        if axis is not None and array.ndim == 0:
+            pytest.skip(f"`axis != None` not supported for 0d arrays")
+
+        array = array.astype(dtype)
+        for func, identity in zip(self.nanfuncs, [0, 1]):
+            out = func(array, axis=axis)
+            assert np.all(out == identity)
+            assert out.dtype == array.dtype
 
     def test_empty(self):
         for f, tgt_value in zip([np.nansum, np.nanprod], [0, 1]):
@@ -451,19 +462,21 @@ class TestNanFunctions_CumSumProd(SharedNanFunctionsTestsMixin):
     nanfuncs = [np.nancumsum, np.nancumprod]
     stdfuncs = [np.cumsum, np.cumprod]
 
-    def test_allnans(self):
-        for f, tgt_value in zip(self.nanfuncs, [0, 1]):
-            # Unlike other nan-functions, sum/prod/cumsum/cumprod don't warn on all nan input
-            with assert_no_warnings():
-                res = f([np.nan]*3, axis=None)
-                tgt = tgt_value*np.ones((3))
-                assert_(np.array_equal(res, tgt), 'result is not %s * np.ones((3))' % (tgt_value))
-                # Check scalar
-                res = f(np.nan)
-                tgt = tgt_value*np.ones((1))
-                assert_(np.array_equal(res, tgt), 'result is not %s * np.ones((1))' % (tgt_value))
-                # Check there is no warning for not all-nan
-                f([0]*3, axis=None)
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"])
+    @pytest.mark.parametrize("array", [
+        np.full((3, 3), np.nan),
+        np.array(np.nan),
+    ], ids=["0d", "2d"])
+    def test_allnans(self, axis, dtype, array):
+        if axis is not None and array.ndim == 0:
+            pytest.skip(f"`axis != None` not supported for 0d arrays")
+
+        array = array.astype(dtype)
+        for func, identity in zip(self.nanfuncs, [0, 1]):
+            out = func(array)
+            assert np.all(out == identity)
+            assert out.dtype == array.dtype
 
     def test_empty(self):
         for f, tgt_value in zip(self.nanfuncs, [0, 1]):
@@ -559,19 +572,29 @@ class TestNanFunctions_MeanVarStd(SharedNanFunctionsTestsMixin):
                     else:
                         assert_(len(sup.log) == 0)
 
-    def test_allnans(self):
-        mat = np.array([np.nan]*9).reshape(3, 3)
-        for f in self.nanfuncs:
-            for axis in [None, 0, 1]:
-                with warnings.catch_warnings(record=True) as w:
-                    warnings.simplefilter('always')
-                    assert_(np.isnan(f(mat, axis=axis)).all())
-                    assert_(len(w) == 1)
-                    assert_(issubclass(w[0].category, RuntimeWarning))
-                    # Check scalar
-                    assert_(np.isnan(f(np.nan)))
-                    assert_(len(w) == 2)
-                    assert_(issubclass(w[0].category, RuntimeWarning))
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"])
+    @pytest.mark.parametrize("array", [
+        np.array(np.nan),
+        np.full((3, 3), np.nan),
+    ], ids=["0d", "2d"])
+    def test_allnans(self, axis, dtype, array):
+        if axis is not None and array.ndim == 0:
+            pytest.skip(f"`axis != None` not supported for 0d arrays")
+
+        array = array.astype(dtype)
+        match = "(Degrees of freedom <= 0 for slice.)|(Mean of empty slice)"
+        for func in self.nanfuncs:
+            with pytest.warns(RuntimeWarning, match=match):
+                out = func(array, axis=axis)
+            assert np.isnan(out).all()
+
+            # `nanvar` and `nanstd` convert complex inputs to their
+            # corresponding floating dtype
+            if func is np.nanmean:
+                assert out.dtype == array.dtype
+            else:
+                assert out.dtype == np.abs(array).dtype
 
     def test_empty(self):
         mat = np.zeros((0, 3))
@@ -835,24 +858,21 @@ class TestNanFunctions_Percentile:
         res = np.nanpercentile(_ndat, (28, 98), axis=1)
         assert_almost_equal(res, tgt)
 
-    def test_allnans(self):
-        mat = np.array([np.nan]*9).reshape(3, 3)
-        for axis in [None, 0, 1]:
-            with warnings.catch_warnings(record=True) as w:
-                warnings.simplefilter('always')
-                assert_(np.isnan(np.nanpercentile(mat, 60, axis=axis)).all())
-                if axis is None:
-                    assert_(len(w) == 1)
-                else:
-                    assert_(len(w) == 3)
-                assert_(issubclass(w[0].category, RuntimeWarning))
-                # Check scalar
-                assert_(np.isnan(np.nanpercentile(np.nan, 60)))
-                if axis is None:
-                    assert_(len(w) == 2)
-                else:
-                    assert_(len(w) == 4)
-                assert_(issubclass(w[0].category, RuntimeWarning))
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"])
+    @pytest.mark.parametrize("array", [
+        np.array(np.nan),
+        np.full((3, 3), np.nan),
+    ], ids=["0d", "2d"])
+    def test_allnans(self, axis, dtype, array):
+        if axis is not None and array.ndim == 0:
+            pytest.skip(f"`axis != None` not supported for 0d arrays")
+
+        array = array.astype(dtype)
+        with pytest.warns(RuntimeWarning, match="All-NaN slice encountered"):
+            out = np.nanpercentile(array, 60, axis=axis)
+        assert np.isnan(out).all()
+        assert out.dtype == array.dtype
 
     def test_empty(self):
         mat = np.zeros((0, 3))
@@ -945,6 +965,22 @@ class TestNanFunctions_Quantile:
         np.nanquantile(np.arange(100.), p, interpolation="midpoint")
         assert_array_equal(p, p0)
 
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"])
+    @pytest.mark.parametrize("array", [
+        np.array(np.nan),
+        np.full((3, 3), np.nan),
+    ], ids=["0d", "2d"])
+    def test_allnans(self, axis, dtype, array):
+        if axis is not None and array.ndim == 0:
+            pytest.skip(f"`axis != None` not supported for 0d arrays")
+
+        array = array.astype(dtype)
+        with pytest.warns(RuntimeWarning, match="All-NaN slice encountered"):
+            out = np.nanquantile(array, 1, axis=axis)
+        assert np.isnan(out).all()
+        assert out.dtype == array.dtype
+
 @pytest.mark.parametrize("arr, expected", [
     # array of floats with some nans
     (np.array([np.nan, 5.0, np.nan, np.inf]),