Merge pull request #18329 from seberg/issue-18286

BUG: Allow unmodified use of isclose, allclose, etc. with timedelta

5 files changed, 36 insertions, 5 deletions

diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index 086439656..89f56fa09 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -2350,8 +2350,13 @@ def isclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
     # Make sure y is an inexact type to avoid bad behavior on abs(MIN_INT).
     # This will cause casting of x later. Also, make sure to allow subclasses
     # (e.g., for numpy.ma).
-    dt = multiarray.result_type(y, 1.)
-    y = array(y, dtype=dt, copy=False, subok=True)
+    # NOTE: We explicitly allow timedelta, which used to work. This could
+    #       possibly be deprecated. See also gh-18286.
+    #       timedelta works if `atol` is an integer or also a timedelta.
+    #       Although, the default tolerances are unlikely to be useful
+    if y.dtype.kind != "m":
+        dt = multiarray.result_type(y, 1.)
+        y = array(y, dtype=dt, copy=False, subok=True)
 
     xfin = isfinite(x)
     yfin = isfinite(y)
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
index f8b388b6f..a697e5faf 100644
--- a/numpy/core/tests/test_numeric.py
+++ b/numpy/core/tests/test_numeric.py
@@ -2578,6 +2578,15 @@ class TestIsclose:
         assert_(np.isclose(0, np.inf) is np.False_)
         assert_(type(np.isclose(0, np.inf)) is np.bool_)
 
+    def test_timedelta(self):
+        # Allclose currently works for timedelta64 as long as `atol` is
+        # an integer or also a timedelta64
+        a = np.array([[1, 2, 3, "NaT"]], dtype="m8[ns]")
+        assert np.isclose(a, a, atol=0, equal_nan=True).all()
+        assert np.isclose(a, a, atol=np.timedelta64(1, "ns"), equal_nan=True).all()
+        assert np.allclose(a, a, atol=0, equal_nan=True)
+        assert np.allclose(a, a, atol=np.timedelta64(1, "ns"), equal_nan=True)
+
 
 class TestStdVar:
     def setup(self):
diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 38a0a8b50..cda2eeb34 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -7869,9 +7869,14 @@ def allclose(a, b, masked_equal=True, rtol=1e-5, atol=1e-8):
 
     # make sure y is an inexact type to avoid abs(MIN_INT); will cause
     # casting of x later.
-    dtype = np.result_type(y, 1.)
-    if y.dtype != dtype:
-        y = masked_array(y, dtype=dtype, copy=False)
+    # NOTE: We explicitly allow timedelta, which used to work. This could
+    #       possibly be deprecated. See also gh-18286.
+    #       timedelta works if `atol` is an integer or also a timedelta.
+    #       Although, the default tolerances are unlikely to be useful
+    if y.dtype.kind != "m":
+        dtype = np.result_type(y, 1.)
+        if y.dtype != dtype:
+            y = masked_array(y, dtype=dtype, copy=False)
 
     m = mask_or(getmask(x), getmask(y))
     xinf = np.isinf(masked_array(x, copy=False, mask=m)).filled(False)
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index 0b2e7303c..f40780625 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -3037,6 +3037,13 @@ class TestMaskedArrayMethods:
         a = masked_array([np.iinfo(np.int_).min], dtype=np.int_)
         assert_(allclose(a, a))
 
+    def test_allclose_timedelta(self):
+        # Allclose currently works for timedelta64 as long as `atol` is
+        # an integer or also a timedelta64
+        a = np.array([[1, 2, 3, 4]], dtype="m8[ns]")
+        assert allclose(a, a, atol=0)
+        assert allclose(a, a, atol=np.timedelta64(1, "ns"))
+
     def test_allany(self):
         # Checks the any/all methods/functions.
         x = np.array([[0.13, 0.26, 0.90],
diff --git a/numpy/testing/tests/test_utils.py b/numpy/testing/tests/test_utils.py
index c3b9e04b6..261ed9705 100644
--- a/numpy/testing/tests/test_utils.py
+++ b/numpy/testing/tests/test_utils.py
@@ -904,6 +904,11 @@ class TestAssertAllclose:
         msg = str(exc_info.value)
         assert_('Max relative difference: 0.5' in msg)
 
+    def test_timedelta(self):
+        # see gh-18286
+        a = np.array([[1, 2, 3, "NaT"]], dtype="m8[ns]")
+        assert_allclose(a, a)
+
 
 class TestArrayAlmostEqualNulp: