diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 253b582eddf2c..9d32bc008eb25 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -1141,6 +1141,7 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray: i8 = self.asi8 result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan) + dtype = tz_to_dtype(tz=other.tz, unit=self._unit) res_values = result.view(f"M8[{self._unit}]") return DatetimeArray._simple_new(res_values, dtype=dtype, freq=self.freq) @@ -1289,13 +1290,21 @@ def _add_timedelta_arraylike( other = ensure_wrapped_if_datetimelike(other) other = cast("TimedeltaArray", other) + self = cast("DatetimeArray | TimedeltaArray", self) + + if self._reso != other._reso: + raise NotImplementedError( + f"Addition of {type(self).__name__} with TimedeltaArray with " + "mis-matched resolutions is not yet supported." + ) self_i8 = self.asi8 other_i8 = other.asi8 new_values = checked_add_with_arr( self_i8, other_i8, arr_mask=self._isnan, b_mask=other._isnan ) - return type(self)(new_values, dtype=self.dtype) + res_values = new_values.view(self._ndarray.dtype) + return type(self)._simple_new(res_values, dtype=self.dtype) @final def _add_nat(self): diff --git a/pandas/core/arrays/timedeltas.py b/pandas/core/arrays/timedeltas.py index 4011f29855949..193ab6dc99350 100644 --- a/pandas/core/arrays/timedeltas.py +++ b/pandas/core/arrays/timedeltas.py @@ -697,9 +697,10 @@ def __rdivmod__(self, other): return res1, res2 def __neg__(self) -> TimedeltaArray: + freq = None if self.freq is not None: - return type(self)(-self._ndarray, freq=-self.freq) - return type(self)(-self._ndarray) + freq = -self.freq + return type(self)._simple_new(-self._ndarray, dtype=self.dtype, freq=freq) def __pos__(self) -> TimedeltaArray: return type(self)(self._ndarray.copy(), freq=self.freq) diff --git a/pandas/tests/arrays/test_timedeltas.py b/pandas/tests/arrays/test_timedeltas.py index b3b79bd988ad8..441b3d6495221 100644 --- a/pandas/tests/arrays/test_timedeltas.py +++ b/pandas/tests/arrays/test_timedeltas.py @@ -167,6 +167,24 @@ def test_div_td_array(self, tda): expected = tda._ndarray / other tm.assert_numpy_array_equal(result, expected) + def test_add_timedeltaarraylike(self, tda): + # TODO(2.0): just do `tda_nano = tda.astype("m8[ns]")` + tda_nano = TimedeltaArray(tda._ndarray.astype("m8[ns]")) + + msg = "mis-matched resolutions is not yet supported" + with pytest.raises(NotImplementedError, match=msg): + tda_nano + tda + with pytest.raises(NotImplementedError, match=msg): + tda + tda_nano + with pytest.raises(NotImplementedError, match=msg): + tda - tda_nano + with pytest.raises(NotImplementedError, match=msg): + tda_nano - tda + + result = tda_nano + tda_nano + expected = tda_nano * 2 + tm.assert_extension_array_equal(result, expected) + class TestTimedeltaArray: @pytest.mark.parametrize("dtype", [int, np.int32, np.int64, "uint32", "uint64"])