Datetimelike __setitem__ (#24477)

Author: TomAugspurger

pandas/core/arrays/datetimelike.py

@@ -478,6 +478,56 @@ def __getitem__(self, key):
 
         return self._simple_new(result, **attribs)
 
+    def __setitem__(
+            self,
+            key,    # type: Union[int, Sequence[int], Sequence[bool], slice]
+            value,  # type: Union[NaTType, Scalar, Sequence[Scalar]]
+    ):
+        # type: (...) -> None
+        # I'm fudging the types a bit here. The "Scalar" above really depends
+        # on type(self). For PeriodArray, it's Period (or stuff coercible
+        # to a period in from_sequence). For DatetimeArray, it's Timestamp...
+        # I don't know if mypy can do that, possibly with Generics.
+        # https://mypy.readthedocs.io/en/latest/generics.html
+
+        if is_list_like(value):
+            is_slice = isinstance(key, slice)
+
+            if lib.is_scalar(key):
+                raise ValueError("setting an array element with a sequence.")
+
+            if (not is_slice
+                    and len(key) != len(value)
+                    and not com.is_bool_indexer(key)):
+                msg = ("shape mismatch: value array of length '{}' does not "
+                       "match indexing result of length '{}'.")
+                raise ValueError(msg.format(len(key), len(value)))
+            if not is_slice and len(key) == 0:
+                return
+
+            value = type(self)._from_sequence(value, dtype=self.dtype)
+            self._check_compatible_with(value)
+            value = value.asi8
+        elif isinstance(value, self._scalar_type):
+            self._check_compatible_with(value)
+            value = self._unbox_scalar(value)
+        elif isna(value) or value == iNaT:
+            value = iNaT
+        else:
+            msg = (
+                "'value' should be a '{scalar}', 'NaT', or array of those. "
+                "Got '{typ}' instead."
+            )
+            raise TypeError(msg.format(scalar=self._scalar_type.__name__,
+                                       typ=type(value).__name__))
+        self._data[key] = value
+        self._maybe_clear_freq()
+
+    def _maybe_clear_freq(self):
+        # inplace operations like __setitem__ may invalidate the freq of
+        # DatetimeArray and TimedeltaArray
+        pass
+
     def astype(self, dtype, copy=True):
         # Some notes on cases we don't have to handle here in the base class:
         #   1. PeriodArray.astype handles period -> period

pandas/core/arrays/datetimes.py

@@ -368,6 +368,9 @@ def _check_compatible_with(self, other):
             raise ValueError("Timezones don't match. '{own} != {other}'"
                              .format(own=self.tz, other=other.tz))
 
+    def _maybe_clear_freq(self):
+        self._freq = None
+
     # -----------------------------------------------------------------
     # Descriptive Properties
 

pandas/core/arrays/period.py

@@ -371,48 +371,6 @@ def _formatter(self, boxed=False):
             return str
         return "'{}'".format
 
-    def __setitem__(
-            self,
-            key,   # type: Union[int, Sequence[int], Sequence[bool], slice]
-            value  # type: Union[NaTType, Period, Sequence[Period]]
-    ):
-        # type: (...) -> None
-        # n.b. the type on `value` is a bit too restrictive.
-        # we also accept a sequence of stuff coercible to a PeriodArray
-        # by period_array, which includes things like ndarray[object],
-        # ndarray[datetime64ns]. I think ndarray[int] / ndarray[str] won't
-        # work, since the freq can't be inferred.
-        if is_list_like(value):
-            is_slice = isinstance(key, slice)
-            if (not is_slice
-                    and len(key) != len(value)
-                    and not com.is_bool_indexer(key)):
-                msg = ("shape mismatch: value array of length '{}' does not "
-                       "match indexing result of length '{}'.")
-                raise ValueError(msg.format(len(key), len(value)))
-            if not is_slice and len(key) == 0:
-                return
-
-            value = period_array(value)
-
-            if self.freqstr != value.freqstr:
-                _raise_on_incompatible(self, value)
-
-            value = value.asi8
-        elif isinstance(value, Period):
-
-            if self.freqstr != value.freqstr:
-                _raise_on_incompatible(self, value)
-
-            value = value.ordinal
-        elif isna(value):
-            value = iNaT
-        else:
-            msg = ("'value' should be a 'Period', 'NaT', or array of those. "
-                   "Got '{}' instead.".format(type(value).__name__))
-            raise TypeError(msg)
-        self._data[key] = value
-
     @Appender(dtl.DatetimeLikeArrayMixin._validate_fill_value.__doc__)
     def _validate_fill_value(self, fill_value):
         if isna(fill_value):

pandas/core/arrays/timedeltas.py

@@ -238,6 +238,9 @@ def _check_compatible_with(self, other):
         # we don't have anything to validate.
         pass
 
+    def _maybe_clear_freq(self):
+        self._freq = None
+
     # ----------------------------------------------------------------
     # Array-Like / EA-Interface Methods
 

pandas/core/indexes/datetimelike.py

@@ -38,6 +38,7 @@ class DatetimeIndexOpsMixin(DatetimeLikeArrayMixin):
     # override DatetimeLikeArrayMixin method
     copy = Index.copy
     view = Index.view
+    __setitem__ = Index.__setitem__
 
     # DatetimeLikeArrayMixin assumes subclasses are mutable, so these are
     # properties there.  They can be made into cache_readonly for Index

pandas/tests/arrays/test_datetimelike.py

@@ -182,6 +182,31 @@ def test_searchsorted(self):
         result = arr.searchsorted(pd.NaT)
         assert result == 0
 
+    def test_setitem(self):
+        data = np.arange(10, dtype='i8') * 24 * 3600 * 10**9
+        arr = self.array_cls(data, freq='D')
+
+        arr[0] = arr[1]
+        expected = np.arange(10, dtype='i8') * 24 * 3600 * 10**9
+        expected[0] = expected[1]
+
+        tm.assert_numpy_array_equal(arr.asi8, expected)
+
+        arr[:2] = arr[-2:]
+        expected[:2] = expected[-2:]
+        tm.assert_numpy_array_equal(arr.asi8, expected)
+
+    def test_setitem_raises(self):
+        data = np.arange(10, dtype='i8') * 24 * 3600 * 10**9
+        arr = self.array_cls(data, freq='D')
+        val = arr[0]
+
+        with pytest.raises(IndexError, match="index 12 is out of bounds"):
+            arr[12] = val
+
+        with pytest.raises(TypeError, match="'value' should be a.* 'object'"):
+            arr[0] = object()
+
 
 class TestDatetimeArray(SharedTests):
     index_cls = pd.DatetimeIndex

pandas/tests/arrays/test_datetimes.py

@@ -74,3 +74,19 @@ def test_tz_setter_raises(self):
         arr = DatetimeArray._from_sequence(['2000'], tz='US/Central')
         with pytest.raises(AttributeError, match='tz_localize'):
             arr.tz = 'UTC'
+
+    def test_setitem_different_tz_raises(self):
+        data = np.array([1, 2, 3], dtype='M8[ns]')
+        arr = DatetimeArray(data, copy=False,
+                            dtype=DatetimeTZDtype(tz="US/Central"))
+        with pytest.raises(ValueError, match="None"):
+            arr[0] = pd.Timestamp('2000')
+
+        with pytest.raises(ValueError, match="US/Central"):
+            arr[0] = pd.Timestamp('2000', tz="US/Eastern")
+
+    def test_setitem_clears_freq(self):
+        a = DatetimeArray(pd.date_range('2000', periods=2, freq='D',
+                                        tz='US/Central'))
+        a[0] = pd.Timestamp("2000", tz="US/Central")
+        assert a.freq is None

pandas/tests/arrays/test_timedeltas.py

@@ -72,3 +72,8 @@ def test_astype_int(self, dtype):
 
         assert result.dtype == expected_dtype
         tm.assert_numpy_array_equal(result, expected)
+
+    def test_setitem_clears_freq(self):
+        a = TimedeltaArray(pd.timedelta_range('1H', periods=2, freq='H'))
+        a[0] = pd.Timedelta("1H")
+        assert a.freq is None

pandas/tests/extension/base/setitem.py

@@ -182,3 +182,8 @@ def test_setitem_slice_array(self, data):
         arr = data[:5].copy()
         arr[:5] = data[-5:]
         self.assert_extension_array_equal(arr, data[-5:])
+
+    def test_setitem_scalar_key_sequence_raise(self, data):
+        arr = data[:5].copy()
+        with pytest.raises(ValueError):
+            arr[0] = arr[[0, 1]]

pandas/tests/extension/decimal/array.py

@@ -108,6 +108,8 @@ def astype(self, dtype, copy=True):
 
     def __setitem__(self, key, value):
         if pd.api.types.is_list_like(value):
+            if pd.api.types.is_scalar(key):
+                raise ValueError("setting an array element with a sequence.")
             value = [decimal.Decimal(v) for v in value]
         else:
             value = decimal.Decimal(value)