BUG: Series[dt64tz].__setitem__(mask, different_tz) raises (#39656)

Author: jbrockmendel

pandas/core/indexes/base.py

@@ -4528,7 +4528,7 @@ def putmask(self, mask, value):
         if not mask.any():
             return self.copy()
 
-        if value is None:
+        if value is None and (self._is_numeric_dtype or self.dtype == object):
             value = self._na_value
         try:
             converted = self._validate_fill_value(value)

pandas/core/indexes/extension.py

@@ -353,11 +353,6 @@ def insert(self: _T, loc: int, item) -> _T:
         new_arr = arr._from_backing_data(new_vals)
         return type(self)._simple_new(new_arr, name=self.name)
 
-    @doc(Index.where)
-    def where(self: _T, cond: np.ndarray, other=None) -> _T:
-        res_values = self._data.where(cond, other)
-        return type(self)._simple_new(res_values, name=self.name)
-
     def putmask(self, mask, value):
         res_values = self._data.copy()
         try:

pandas/core/internals/blocks.py

@@ -1014,14 +1014,6 @@ def putmask(self, mask, new) -> List[Block]:
             new = self.fill_value
 
         if self._can_hold_element(new):
-            if self.dtype.kind in ["m", "M"]:
-                arr = self.array_values()
-                arr = cast("NDArrayBackedExtensionArray", arr)
-                if transpose:
-                    arr = arr.T
-                arr.putmask(mask, new)
-                return [self]
-
             if transpose:
                 new_values = new_values.T
 
@@ -2025,6 +2017,18 @@ def to_native_types(self, na_rep="NaT", **kwargs):
         result = arr._format_native_types(na_rep=na_rep, **kwargs)
         return self.make_block(result)
 
+    def putmask(self, mask, new) -> List[Block]:
+        mask = _extract_bool_array(mask)
+
+        if not self._can_hold_element(new):
+            return self.astype(object).putmask(mask, new)
+
+        # TODO(EA2D): reshape unnecessary with 2D EAs
+        arr = self.array_values().reshape(self.shape)
+        arr = cast("NDArrayBackedExtensionArray", arr)
+        arr.T.putmask(mask, new)
+        return [self]
+
     def where(self, other, cond, errors="raise", axis: int = 0) -> List[Block]:
         # TODO(EA2D): reshape unnecessary with 2D EAs
         arr = self.array_values().reshape(self.shape)
@@ -2099,6 +2103,7 @@ class DatetimeTZBlock(ExtensionBlock, DatetimeBlock):
     diff = DatetimeBlock.diff
     fill_value = np.datetime64("NaT", "ns")
     where = DatetimeBlock.where
+    putmask = DatetimeLikeBlockMixin.putmask
 
     array_values = ExtensionBlock.array_values
 

pandas/tests/indexes/categorical/test_indexing.py

@@ -305,10 +305,11 @@ def test_where_non_categories(self):
         ci = CategoricalIndex(["a", "b", "c", "d"])
         mask = np.array([True, False, True, False])
 
-        msg = "Cannot setitem on a Categorical with a new category"
-        with pytest.raises(ValueError, match=msg):
-            ci.where(mask, 2)
+        result = ci.where(mask, 2)
+        expected = Index(["a", 2, "c", 2], dtype=object)
+        tm.assert_index_equal(result, expected)
 
+        msg = "Cannot setitem on a Categorical with a new category"
         with pytest.raises(ValueError, match=msg):
             # Test the Categorical method directly
             ci._data.where(mask, 2)

pandas/tests/indexes/datetimelike.py

@@ -117,9 +117,9 @@ def test_where_cast_str(self):
         result = index.where(mask, [str(index[0])])
         tm.assert_index_equal(result, expected)
 
-        msg = "value should be a '.*', 'NaT', or array of those"
-        with pytest.raises(TypeError, match=msg):
-            index.where(mask, "foo")
+        expected = index.astype(object).where(mask, "foo")
+        result = index.where(mask, "foo")
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            index.where(mask, ["foo"])
+        result = index.where(mask, ["foo"])
+        tm.assert_index_equal(result, expected)

pandas/tests/indexes/datetimes/test_indexing.py

@@ -175,40 +175,56 @@ def test_where_other(self):
     def test_where_invalid_dtypes(self):
         dti = date_range("20130101", periods=3, tz="US/Eastern")
 
-        i2 = Index([pd.NaT, pd.NaT] + dti[2:].tolist())
+        tail = dti[2:].tolist()
+        i2 = Index([pd.NaT, pd.NaT] + tail)
 
-        msg = "value should be a 'Timestamp', 'NaT', or array of those. Got"
-        msg2 = "Cannot compare tz-naive and tz-aware datetime-like objects"
-        with pytest.raises(TypeError, match=msg2):
-            # passing tz-naive ndarray to tzaware DTI
-            dti.where(notna(i2), i2.values)
+        mask = notna(i2)
 
-        with pytest.raises(TypeError, match=msg2):
-            # passing tz-aware DTI to tznaive DTI
-            dti.tz_localize(None).where(notna(i2), i2)
+        # passing tz-naive ndarray to tzaware DTI
+        result = dti.where(mask, i2.values)
+        expected = Index([pd.NaT.asm8, pd.NaT.asm8] + tail, dtype=object)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            dti.where(notna(i2), i2.tz_localize(None).to_period("D"))
+        # passing tz-aware DTI to tznaive DTI
+        naive = dti.tz_localize(None)
+        result = naive.where(mask, i2)
+        expected = Index([i2[0], i2[1]] + naive[2:].tolist(), dtype=object)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            dti.where(notna(i2), i2.asi8.view("timedelta64[ns]"))
+        pi = i2.tz_localize(None).to_period("D")
+        result = dti.where(mask, pi)
+        expected = Index([pi[0], pi[1]] + tail, dtype=object)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            dti.where(notna(i2), i2.asi8)
+        tda = i2.asi8.view("timedelta64[ns]")
+        result = dti.where(mask, tda)
+        expected = Index([tda[0], tda[1]] + tail, dtype=object)
+        assert isinstance(expected[0], np.timedelta64)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            # non-matching scalar
-            dti.where(notna(i2), pd.Timedelta(days=4))
+        result = dti.where(mask, i2.asi8)
+        expected = Index([pd.NaT.value, pd.NaT.value] + tail, dtype=object)
+        assert isinstance(expected[0], int)
+        tm.assert_index_equal(result, expected)
+
+        # non-matching scalar
+        td = pd.Timedelta(days=4)
+        result = dti.where(mask, td)
+        expected = Index([td, td] + tail, dtype=object)
+        assert expected[0] is td
+        tm.assert_index_equal(result, expected)
 
     def test_where_mismatched_nat(self, tz_aware_fixture):
         tz = tz_aware_fixture
         dti = date_range("2013-01-01", periods=3, tz=tz)
         cond = np.array([True, False, True])
 
-        msg = "value should be a 'Timestamp', 'NaT', or array of those. Got"
-        with pytest.raises(TypeError, match=msg):
-            # wrong-dtyped NaT
-            dti.where(cond, np.timedelta64("NaT", "ns"))
+        tdnat = np.timedelta64("NaT", "ns")
+        expected = Index([dti[0], tdnat, dti[2]], dtype=object)
+        assert expected[1] is tdnat
+
+        result = dti.where(cond, tdnat)
+        tm.assert_index_equal(result, expected)
 
     def test_where_tz(self):
         i = date_range("20130101", periods=3, tz="US/Eastern")

pandas/tests/indexes/period/test_indexing.py

@@ -603,30 +603,42 @@ def test_where_other(self):
     def test_where_invalid_dtypes(self):
         pi = period_range("20130101", periods=5, freq="D")
 
-        i2 = PeriodIndex([NaT, NaT] + pi[2:].tolist(), freq="D")
+        tail = pi[2:].tolist()
+        i2 = PeriodIndex([NaT, NaT] + tail, freq="D")
+        mask = notna(i2)
 
-        msg = "value should be a 'Period', 'NaT', or array of those"
-        with pytest.raises(TypeError, match=msg):
-            pi.where(notna(i2), i2.asi8)
+        result = pi.where(mask, i2.asi8)
+        expected = pd.Index([NaT.value, NaT.value] + tail, dtype=object)
+        assert isinstance(expected[0], int)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            pi.where(notna(i2), i2.asi8.view("timedelta64[ns]"))
+        tdi = i2.asi8.view("timedelta64[ns]")
+        expected = pd.Index([tdi[0], tdi[1]] + tail, dtype=object)
+        assert isinstance(expected[0], np.timedelta64)
+        result = pi.where(mask, tdi)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            pi.where(notna(i2), i2.to_timestamp("S"))
+        dti = i2.to_timestamp("S")
+        expected = pd.Index([dti[0], dti[1]] + tail, dtype=object)
+        assert expected[0] is NaT
+        result = pi.where(mask, dti)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            # non-matching scalar
-            pi.where(notna(i2), Timedelta(days=4))
+        td = Timedelta(days=4)
+        expected = pd.Index([td, td] + tail, dtype=object)
+        assert expected[0] == td
+        result = pi.where(mask, td)
+        tm.assert_index_equal(result, expected)
 
     def test_where_mismatched_nat(self):
         pi = period_range("20130101", periods=5, freq="D")
         cond = np.array([True, False, True, True, False])
 
-        msg = "value should be a 'Period', 'NaT', or array of those"
-        with pytest.raises(TypeError, match=msg):
-            # wrong-dtyped NaT
-            pi.where(cond, np.timedelta64("NaT", "ns"))
+        tdnat = np.timedelta64("NaT", "ns")
+        expected = pd.Index([pi[0], tdnat, pi[2], pi[3], tdnat], dtype=object)
+        assert expected[1] is tdnat
+        result = pi.where(cond, tdnat)
+        tm.assert_index_equal(result, expected)
 
 
 class TestTake:

pandas/tests/indexes/timedeltas/test_indexing.py

@@ -148,30 +148,39 @@ def test_where_doesnt_retain_freq(self):
     def test_where_invalid_dtypes(self):
         tdi = timedelta_range("1 day", periods=3, freq="D", name="idx")
 
-        i2 = Index([pd.NaT, pd.NaT] + tdi[2:].tolist())
+        tail = tdi[2:].tolist()
+        i2 = Index([pd.NaT, pd.NaT] + tail)
+        mask = notna(i2)
 
-        msg = "value should be a 'Timedelta', 'NaT', or array of those"
-        with pytest.raises(TypeError, match=msg):
-            tdi.where(notna(i2), i2.asi8)
+        expected = Index([pd.NaT.value, pd.NaT.value] + tail, dtype=object, name="idx")
+        assert isinstance(expected[0], int)
+        result = tdi.where(mask, i2.asi8)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            tdi.where(notna(i2), i2 + pd.Timestamp.now())
+        ts = i2 + pd.Timestamp.now()
+        expected = Index([ts[0], ts[1]] + tail, dtype=object, name="idx")
+        result = tdi.where(mask, ts)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            tdi.where(notna(i2), (i2 + pd.Timestamp.now()).to_period("D"))
+        per = (i2 + pd.Timestamp.now()).to_period("D")
+        expected = Index([per[0], per[1]] + tail, dtype=object, name="idx")
+        result = tdi.where(mask, per)
+        tm.assert_index_equal(result, expected)
 
-        with pytest.raises(TypeError, match=msg):
-            # non-matching scalar
-            tdi.where(notna(i2), pd.Timestamp.now())
+        ts = pd.Timestamp.now()
+        expected = Index([ts, ts] + tail, dtype=object, name="idx")
+        result = tdi.where(mask, ts)
+        tm.assert_index_equal(result, expected)
 
     def test_where_mismatched_nat(self):
         tdi = timedelta_range("1 day", periods=3, freq="D", name="idx")
         cond = np.array([True, False, False])
 
-        msg = "value should be a 'Timedelta', 'NaT', or array of those"
-        with pytest.raises(TypeError, match=msg):
-            # wrong-dtyped NaT
-            tdi.where(cond, np.datetime64("NaT", "ns"))
+        dtnat = np.datetime64("NaT", "ns")
+        expected = Index([tdi[0], dtnat, dtnat], dtype=object, name="idx")
+        assert expected[2] is dtnat
+        result = tdi.where(cond, dtnat)
+        tm.assert_index_equal(result, expected)
 
 
 class TestTake:

pandas/tests/indexing/test_coercion.py

@@ -802,10 +802,13 @@ def test_where_index_timedelta64(self, value):
         result = tdi.where(cond, value)
         tm.assert_index_equal(result, expected)
 
-        msg = "value should be a 'Timedelta', 'NaT', or array of thos"
-        with pytest.raises(TypeError, match=msg):
-            # wrong-dtyped NaT
-            tdi.where(cond, np.datetime64("NaT", "ns"))
+        # wrong-dtyped NaT
+        dtnat = np.datetime64("NaT", "ns")
+        expected = pd.Index([tdi[0], dtnat, dtnat, tdi[3]], dtype=object)
+        assert expected[1] is dtnat
+
+        result = tdi.where(cond, dtnat)
+        tm.assert_index_equal(result, expected)
 
     def test_where_index_period(self):
         dti = pd.date_range("2016-01-01", periods=3, freq="QS")
@@ -825,14 +828,16 @@ def test_where_index_period(self):
         expected = pd.PeriodIndex([other[0], pi[1], other[2]])
         tm.assert_index_equal(result, expected)
 
-        # Passing a mismatched scalar
-        msg = "value should be a 'Period', 'NaT', or array of those"
-        with pytest.raises(TypeError, match=msg):
-            pi.where(cond, pd.Timedelta(days=4))
+        # Passing a mismatched scalar -> casts to object
+        td = pd.Timedelta(days=4)
+        expected = pd.Index([td, pi[1], td], dtype=object)
+        result = pi.where(cond, td)
+        tm.assert_index_equal(result, expected)
 
-        msg = r"Input has different freq=D from PeriodArray\(freq=Q-DEC\)"
-        with pytest.raises(ValueError, match=msg):
-            pi.where(cond, pd.Period("2020-04-21", "D"))
+        per = pd.Period("2020-04-21", "D")
+        expected = pd.Index([per, pi[1], per], dtype=object)
+        result = pi.where(cond, per)
+        tm.assert_index_equal(result, expected)
 
 
 class TestFillnaSeriesCoercion(CoercionBase):

pandas/tests/series/indexing/test_setitem.py

@@ -50,19 +50,6 @@ def test_setitem_with_string_index(self):
         assert ser.Date == date.today()
         assert ser["Date"] == date.today()
 
-    def test_setitem_with_different_tz_casts_to_object(self):
-        # GH#24024
-        ser = Series(date_range("2000", periods=2, tz="US/Central"))
-        ser[0] = Timestamp("2000", tz="US/Eastern")
-        expected = Series(
-            [
-                Timestamp("2000-01-01 00:00:00-05:00", tz="US/Eastern"),
-                Timestamp("2000-01-02 00:00:00-06:00", tz="US/Central"),
-            ],
-            dtype=object,
-        )
-        tm.assert_series_equal(ser, expected)
-
     def test_setitem_tuple_with_datetimetz_values(self):
         # GH#20441
         arr = date_range("2017", periods=4, tz="US/Eastern")
@@ -646,3 +633,29 @@ def val(self, request):
     @pytest.fixture
     def is_inplace(self):
         return True
+
+
+class TestSetitemMismatchedTZCastsToObject(SetitemCastingEquivalents):
+    # GH#24024
+    @pytest.fixture
+    def obj(self):
+        return Series(date_range("2000", periods=2, tz="US/Central"))
+
+    @pytest.fixture
+    def val(self):
+        return Timestamp("2000", tz="US/Eastern")
+
+    @pytest.fixture
+    def key(self):
+        return 0
+
+    @pytest.fixture
+    def expected(self):
+        expected = Series(
+            [
+                Timestamp("2000-01-01 00:00:00-05:00", tz="US/Eastern"),
+                Timestamp("2000-01-02 00:00:00-06:00", tz="US/Central"),
+            ],
+            dtype=object,
+        )
+        return expected