BUG: DTI/TDI/PI.where accepting incorrectly-typed NaTs (pandas-dev#33715)

Author: jbrockmendel

pandas/core/arrays/datetimelike.py

@@ -869,27 +869,33 @@ def _validate_insert_value(self, value):
 
     def _validate_where_value(self, other):
         if is_valid_nat_for_dtype(other, self.dtype):
-            other = NaT.value
+            other = NaT
+        elif isinstance(other, self._recognized_scalars):
+            other = self._scalar_type(other)
+            self._check_compatible_with(other, setitem=True)
         elif not is_list_like(other):
-            # TODO: what about own-type scalars?
             raise TypeError(f"Where requires matching dtype, not {type(other)}")
 
         else:
             # Do type inference if necessary up front
             # e.g. we passed PeriodIndex.values and got an ndarray of Periods
-            from pandas import Index
-
-            other = Index(other)
+            other = array(other)
+            other = extract_array(other, extract_numpy=True)
 
-            if is_categorical_dtype(other):
+            if is_categorical_dtype(other.dtype):
                 # e.g. we have a Categorical holding self.dtype
                 if is_dtype_equal(other.categories.dtype, self.dtype):
                     other = other._internal_get_values()
 
-            if not is_dtype_equal(self.dtype, other.dtype):
+            if not type(self)._is_recognized_dtype(other.dtype):
                 raise TypeError(f"Where requires matching dtype, not {other.dtype}")
+            self._check_compatible_with(other, setitem=True)
 
+        if lib.is_scalar(other):
+            other = self._unbox_scalar(other)
+        else:
             other = other.view("i8")
+
         return other
 
     # ------------------------------------------------------------------

pandas/core/indexes/datetimelike.py

@@ -530,7 +530,12 @@ def isin(self, values, level=None):
     def where(self, cond, other=None):
         values = self.view("i8")
 
-        other = self._data._validate_where_value(other)
+        try:
+            other = self._data._validate_where_value(other)
+        except (TypeError, ValueError) as err:
+            # Includes tzawareness mismatch and IncompatibleFrequencyError
+            oth = getattr(other, "dtype", other)
+            raise TypeError(f"Where requires matching dtype, not {oth}") from err
 
         result = np.where(cond, values, other).astype("i8")
         arr = type(self._data)._simple_new(result, dtype=self.dtype)

pandas/tests/indexes/datetimes/test_indexing.py

@@ -197,9 +197,15 @@ def test_where_invalid_dtypes(self):
             # non-matching scalar
             dti.where(notna(i2), pd.Timedelta(days=4))
 
-        with pytest.raises(TypeError, match="Where requires matching dtype"):
-            # non-matching NA value
-            dti.where(notna(i2), np.timedelta64("NaT", "ns"))
+    def test_where_mismatched_nat(self, tz_aware_fixture):
+        tz = tz_aware_fixture
+        dti = pd.date_range("2013-01-01", periods=3, tz=tz)
+        cond = np.array([True, False, True])
+
+        msg = "Where requires matching dtype"
+        with pytest.raises(TypeError, match=msg):
+            # wrong-dtyped NaT
+            dti.where(cond, np.timedelta64("NaT", "ns"))
 
     def test_where_tz(self):
         i = pd.date_range("20130101", periods=3, tz="US/Eastern")

pandas/tests/indexes/period/test_indexing.py

@@ -541,9 +541,14 @@ def test_where_invalid_dtypes(self):
             # non-matching scalar
             pi.where(notna(i2), Timedelta(days=4))
 
-        with pytest.raises(TypeError, match="Where requires matching dtype"):
-            # non-matching NA value
-            pi.where(notna(i2), np.timedelta64("NaT", "ns"))
+    def test_where_mismatched_nat(self):
+        pi = period_range("20130101", periods=5, freq="D")
+        cond = np.array([True, False, True, True, False])
+
+        msg = "Where requires matching dtype"
+        with pytest.raises(TypeError, match=msg):
+            # wrong-dtyped NaT
+            pi.where(cond, np.timedelta64("NaT", "ns"))
 
 
 class TestTake:

pandas/tests/indexes/timedeltas/test_indexing.py

@@ -163,9 +163,14 @@ def test_where_invalid_dtypes(self):
             # non-matching scalar
             tdi.where(notna(i2), pd.Timestamp.now())
 
-        with pytest.raises(TypeError, match="Where requires matching dtype"):
-            # non-matching NA value
-            tdi.where(notna(i2), np.datetime64("NaT", "ns"))
+    def test_where_mismatched_nat(self):
+        tdi = timedelta_range("1 day", periods=3, freq="D", name="idx")
+        cond = np.array([True, False, False])
+
+        msg = "Where requires matching dtype"
+        with pytest.raises(TypeError, match=msg):
+            # wrong-dtyped NaT
+            tdi.where(cond, np.datetime64("NaT", "ns"))
 
 
 class TestTake:

pandas/tests/indexing/test_coercion.py

@@ -1,3 +1,4 @@
+from datetime import timedelta
 import itertools
 from typing import Dict, List
 
@@ -686,8 +687,15 @@ def test_where_series_datetime64(self, fill_val, exp_dtype):
         )
         self._assert_where_conversion(obj, cond, values, exp, exp_dtype)
 
-    def test_where_index_datetime(self):
-        fill_val = pd.Timestamp("2012-01-01")
+    @pytest.mark.parametrize(
+        "fill_val",
+        [
+            pd.Timestamp("2012-01-01"),
+            pd.Timestamp("2012-01-01").to_datetime64(),
+            pd.Timestamp("2012-01-01").to_pydatetime(),
+        ],
+    )
+    def test_where_index_datetime(self, fill_val):
         exp_dtype = "datetime64[ns]"
         obj = pd.Index(
             [
@@ -700,9 +708,9 @@ def test_where_index_datetime(self):
         assert obj.dtype == "datetime64[ns]"
         cond = pd.Index([True, False, True, False])
 
-        msg = "Where requires matching dtype, not .*Timestamp"
-        with pytest.raises(TypeError, match=msg):
-            obj.where(cond, fill_val)
+        result = obj.where(cond, fill_val)
+        expected = pd.DatetimeIndex([obj[0], fill_val, obj[2], fill_val])
+        tm.assert_index_equal(result, expected)
 
         values = pd.Index(pd.date_range(fill_val, periods=4))
         exp = pd.Index(
@@ -717,7 +725,7 @@ def test_where_index_datetime(self):
         self._assert_where_conversion(obj, cond, values, exp, exp_dtype)
 
     @pytest.mark.xfail(reason="GH 22839: do not ignore timezone, must be object")
-    def test_where_index_datetimetz(self):
+    def test_where_index_datetime64tz(self):
         fill_val = pd.Timestamp("2012-01-01", tz="US/Eastern")
         exp_dtype = np.object
         obj = pd.Index(
@@ -754,23 +762,53 @@ def test_where_index_complex128(self):
     def test_where_index_bool(self):
         pass
 
-    def test_where_series_datetime64tz(self):
-        pass
-
     def test_where_series_timedelta64(self):
         pass
 
     def test_where_series_period(self):
         pass
 
-    def test_where_index_datetime64tz(self):
-        pass
+    @pytest.mark.parametrize(
+        "value", [pd.Timedelta(days=9), timedelta(days=9), np.timedelta64(9, "D")]
+    )
+    def test_where_index_timedelta64(self, value):
+        tdi = pd.timedelta_range("1 Day", periods=4)
+        cond = np.array([True, False, False, True])
 
-    def test_where_index_timedelta64(self):
-        pass
+        expected = pd.TimedeltaIndex(["1 Day", value, value, "4 Days"])
+        result = tdi.where(cond, value)
+        tm.assert_index_equal(result, expected)
+
+        msg = "Where requires matching dtype"
+        with pytest.raises(TypeError, match=msg):
+            # wrong-dtyped NaT
+            tdi.where(cond, np.datetime64("NaT", "ns"))
 
     def test_where_index_period(self):
-        pass
+        dti = pd.date_range("2016-01-01", periods=3, freq="QS")
+        pi = dti.to_period("Q")
+
+        cond = np.array([False, True, False])
+
+        # Passinga  valid scalar
+        value = pi[-1] + pi.freq * 10
+        expected = pd.PeriodIndex([value, pi[1], value])
+        result = pi.where(cond, value)
+        tm.assert_index_equal(result, expected)
+
+        # Case passing ndarray[object] of Periods
+        other = np.asarray(pi + pi.freq * 10, dtype=object)
+        result = pi.where(cond, other)
+        expected = pd.PeriodIndex([other[0], pi[1], other[2]])
+        tm.assert_index_equal(result, expected)
+
+        # Passing a mismatched scalar
+        msg = "Where requires matching dtype"
+        with pytest.raises(TypeError, match=msg):
+            pi.where(cond, pd.Timedelta(days=4))
+
+        with pytest.raises(TypeError, match=msg):
+            pi.where(cond, pd.Period("2020-04-21", "D"))
 
 
 class TestFillnaSeriesCoercion(CoercionBase):