From 770b533d3e4e09e6e191ed609cab89049d7a07a2 Mon Sep 17 00:00:00 2001 From: Brock Date: Fri, 30 Sep 2022 11:20:20 -0700 Subject: [PATCH] BUG: DatetimeArray-datetimelike mixed resos --- pandas/_libs/tslibs/timestamps.pyi | 1 + pandas/core/arrays/datetimelike.py | 20 ++++++++++++++++- pandas/tests/arrays/test_datetimes.py | 31 ++++++++++++++++++++++++++- 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/pandas/_libs/tslibs/timestamps.pyi b/pandas/_libs/tslibs/timestamps.pyi index e4be7fda43005..8382fe0274138 100644 --- a/pandas/_libs/tslibs/timestamps.pyi +++ b/pandas/_libs/tslibs/timestamps.pyi @@ -27,6 +27,7 @@ _DatetimeT = TypeVar("_DatetimeT", bound=datetime) def integer_op_not_supported(obj: object) -> TypeError: ... class Timestamp(datetime): + _reso: int min: ClassVar[Timestamp] max: ClassVar[Timestamp] diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 71f834e0f4e01..c42eeb9435c89 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -1218,6 +1218,13 @@ def _sub_datetimelike_scalar(self, other: datetime | np.datetime64): new_message = str(err).replace("compare", "subtract") raise type(err)(new_message) from err + if other._reso != self._reso: + if other._reso < self._reso: + other = other._as_unit(self._unit) + else: + unit = npy_unit_to_abbrev(other._reso) + self = self._as_unit(unit) + i8 = self.asi8 result = checked_add_with_arr(i8, -other.value, arr_mask=self._isnan) res_m8 = result.view(f"timedelta64[{self._unit}]") @@ -1248,12 +1255,23 @@ def _sub_datetime_arraylike(self, other): new_message = str(err).replace("compare", "subtract") raise type(err)(new_message) from err + if other._reso != self._reso: + if other._reso < self._reso: + other = other._as_unit(self._unit) + else: + self = self._as_unit(other._unit) + 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 new_values.view("timedelta64[ns]") + res_m8 = new_values.view(f"timedelta64[{self._unit}]") + + from pandas.core.arrays import TimedeltaArray + + new_freq = self._get_arithmetic_result_freq(other) + return TimedeltaArray._simple_new(res_m8, dtype=res_m8.dtype, freq=new_freq) @final def _sub_period(self, other: Period) -> npt.NDArray[np.object_]: diff --git a/pandas/tests/arrays/test_datetimes.py b/pandas/tests/arrays/test_datetimes.py index 83850c6cd2594..8c2a8df7c2fbd 100644 --- a/pandas/tests/arrays/test_datetimes.py +++ b/pandas/tests/arrays/test_datetimes.py @@ -13,7 +13,10 @@ import pandas as pd import pandas._testing as tm -from pandas.core.arrays import DatetimeArray +from pandas.core.arrays import ( + DatetimeArray, + TimedeltaArray, +) class TestNonNano: @@ -218,6 +221,32 @@ def test_add_mismatched_reso_doesnt_downcast(self): # (so we _could_ downcast to unit="s"), we do not. assert res._unit == "us" + def test_sub_datetimelike_scalar_mismatch(self): + dti = pd.date_range("2016-01-01", periods=3) + dta = dti._data._as_unit("us") + + ts = dta[0]._as_unit("s") + + result = dta - ts + expected = (dti - dti[0])._data._as_unit("us") + assert result.dtype == "m8[us]" + tm.assert_extension_array_equal(result, expected) + + def test_sub_datetime64_reso_mismatch(self): + dti = pd.date_range("2016-01-01", periods=3) + left = dti._data._as_unit("s") + right = left._as_unit("ms") + + result = left - right + exp_values = np.array([0, 0, 0], dtype="m8[ms]") + expected = TimedeltaArray._simple_new( + exp_values, + dtype=exp_values.dtype, + ) + tm.assert_extension_array_equal(result, expected) + result2 = right - left + tm.assert_extension_array_equal(result2, expected) + class TestDatetimeArrayComparisons: # TODO: merge this into tests/arithmetic/test_datetime64 once it is