1 files changed, 55 insertions, 0 deletions
diff --git a/numpy/lib/tests/test_nanfunctions.py b/numpy/lib/tests/test_nanfunctions.py
index 504372faf..da2d0cc52 100644
--- a/numpy/lib/tests/test_nanfunctions.py
+++ b/numpy/lib/tests/test_nanfunctions.py
@@ -1,8 +1,10 @@
 from __future__ import division, absolute_import, print_function
 
 import warnings
+import pytest
 
 import numpy as np
+from numpy.lib.nanfunctions import _nan_mask, _replace_nan
 from numpy.testing import (
     assert_, assert_equal, assert_almost_equal, assert_no_warnings,
     assert_raises, assert_array_equal, suppress_warnings
@@ -925,3 +927,56 @@ class TestNanFunctions_Quantile(object):
         p = p.tolist()
         np.nanquantile(np.arange(100.), p, interpolation="midpoint")
         assert_array_equal(p, p0)
+
+@pytest.mark.parametrize("arr, expected", [
+    # array of floats with some nans
+    (np.array([np.nan, 5.0, np.nan, np.inf]),
+     np.array([False, True, False, True])),
+    # int64 array that can't possibly have nans
+    (np.array([1, 5, 7, 9], dtype=np.int64),
+     True),
+    # bool array that can't possibly have nans
+    (np.array([False, True, False, True]),
+     True),
+    # 2-D complex array with nans
+    (np.array([[np.nan, 5.0],
+               [np.nan, np.inf]], dtype=np.complex64),
+     np.array([[False, True],
+               [False, True]])),
+    ])
+def test__nan_mask(arr, expected):
+    for out in [None, np.empty(arr.shape, dtype=np.bool_)]:
+        actual = _nan_mask(arr, out=out)
+        assert_equal(actual, expected)
+        # the above won't distinguish between True proper
+        # and an array of True values; we want True proper
+        # for types that can't possibly contain NaN
+        if type(expected) is not np.ndarray:
+            assert actual is True
+
+
+def test__replace_nan():
+    """ Test that _replace_nan returns the original array if there are no
+    NaNs, not a copy.
+    """
+    for dtype in [np.bool, np.int32, np.int64]:
+        arr = np.array([0, 1], dtype=dtype)
+        result, mask = _replace_nan(arr, 0)
+        assert mask is None
+        # do not make a copy if there are no nans
+        assert result is arr
+
+    for dtype in [np.float32, np.float64]:
+        arr = np.array([0, 1], dtype=dtype)
+        result, mask = _replace_nan(arr, 2)
+        assert (mask == False).all()
+        # mask is not None, so we make a copy
+        assert result is not arr
+        assert_equal(result, arr)
+
+        arr_nan = np.array([0, 1, np.nan], dtype=dtype)
+        result_nan, mask_nan = _replace_nan(arr_nan, 2)
+        assert_equal(mask_nan, np.array([False, False, True]))
+        assert result_nan is not arr_nan
+        assert_equal(result_nan, np.array([0, 1, 2]))
+        assert np.isnan(arr_nan[-1])