Skip to content

Commit b8dcd27

Browse files
authored
ENH: TimedeltaArray add/sub with non-nano (#48261)
* ENH: TimedeltaArray add/sub with non-nano * mypy fixup
1 parent 03d8cd2 commit b8dcd27

File tree

3 files changed

+31
-3
lines changed

3 files changed

+31
-3
lines changed

pandas/core/arrays/datetimelike.py

+10-1
Original file line numberDiff line numberDiff line change
@@ -1141,6 +1141,7 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray:
11411141

11421142
i8 = self.asi8
11431143
result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan)
1144+
11441145
dtype = tz_to_dtype(tz=other.tz, unit=self._unit)
11451146
res_values = result.view(f"M8[{self._unit}]")
11461147
return DatetimeArray._simple_new(res_values, dtype=dtype, freq=self.freq)
@@ -1289,13 +1290,21 @@ def _add_timedelta_arraylike(
12891290

12901291
other = ensure_wrapped_if_datetimelike(other)
12911292
other = cast("TimedeltaArray", other)
1293+
self = cast("DatetimeArray | TimedeltaArray", self)
1294+
1295+
if self._reso != other._reso:
1296+
raise NotImplementedError(
1297+
f"Addition of {type(self).__name__} with TimedeltaArray with "
1298+
"mis-matched resolutions is not yet supported."
1299+
)
12921300

12931301
self_i8 = self.asi8
12941302
other_i8 = other.asi8
12951303
new_values = checked_add_with_arr(
12961304
self_i8, other_i8, arr_mask=self._isnan, b_mask=other._isnan
12971305
)
1298-
return type(self)(new_values, dtype=self.dtype)
1306+
res_values = new_values.view(self._ndarray.dtype)
1307+
return type(self)._simple_new(res_values, dtype=self.dtype)
12991308

13001309
@final
13011310
def _add_nat(self):

pandas/core/arrays/timedeltas.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -697,9 +697,10 @@ def __rdivmod__(self, other):
697697
return res1, res2
698698

699699
def __neg__(self) -> TimedeltaArray:
700+
freq = None
700701
if self.freq is not None:
701-
return type(self)(-self._ndarray, freq=-self.freq)
702-
return type(self)(-self._ndarray)
702+
freq = -self.freq
703+
return type(self)._simple_new(-self._ndarray, dtype=self.dtype, freq=freq)
703704

704705
def __pos__(self) -> TimedeltaArray:
705706
return type(self)(self._ndarray.copy(), freq=self.freq)

pandas/tests/arrays/test_timedeltas.py

+18
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,24 @@ def test_div_td_array(self, tda):
174174
expected = tda._ndarray / other
175175
tm.assert_numpy_array_equal(result, expected)
176176

177+
def test_add_timedeltaarraylike(self, tda):
178+
# TODO(2.0): just do `tda_nano = tda.astype("m8[ns]")`
179+
tda_nano = TimedeltaArray(tda._ndarray.astype("m8[ns]"))
180+
181+
msg = "mis-matched resolutions is not yet supported"
182+
with pytest.raises(NotImplementedError, match=msg):
183+
tda_nano + tda
184+
with pytest.raises(NotImplementedError, match=msg):
185+
tda + tda_nano
186+
with pytest.raises(NotImplementedError, match=msg):
187+
tda - tda_nano
188+
with pytest.raises(NotImplementedError, match=msg):
189+
tda_nano - tda
190+
191+
result = tda_nano + tda_nano
192+
expected = tda_nano * 2
193+
tm.assert_extension_array_equal(result, expected)
194+
177195

178196
class TestTimedeltaArray:
179197
@pytest.mark.parametrize("dtype", [int, np.int32, np.int64, "uint32", "uint64"])

0 commit comments

Comments
 (0)