diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index a2251c49a2cc5..eadf47b36d7fc 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -73,7 +73,6 @@ from pandas.util._exceptions import find_stack_level from pandas.core.dtypes.common import ( - DT64NS_DTYPE, is_all_strings, is_categorical_dtype, is_datetime64_any_dtype, @@ -1103,6 +1102,7 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray: self = cast("TimedeltaArray", self) from pandas.core.arrays import DatetimeArray + from pandas.core.arrays.datetimes import tz_to_dtype assert other is not NaT other = Timestamp(other) @@ -1113,10 +1113,17 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray: # Preserve our resolution return DatetimeArray._simple_new(result, dtype=result.dtype) + if self._reso != other._reso: + raise NotImplementedError( + "Addition between TimedeltaArray and Timestamp with mis-matched " + "resolutions is not yet supported." + ) + i8 = self.asi8 result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan) - dtype = DatetimeTZDtype(tz=other.tz) if other.tz else DT64NS_DTYPE - return DatetimeArray(result, dtype=dtype, freq=self.freq) + 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) @final def _add_datetime_arraylike(self, other) -> DatetimeArray: diff --git a/pandas/core/arrays/datetimes.py b/pandas/core/arrays/datetimes.py index c9f5946c30c8c..106afcc3c12ea 100644 --- a/pandas/core/arrays/datetimes.py +++ b/pandas/core/arrays/datetimes.py @@ -91,22 +91,23 @@ _midnight = time(0, 0) -def tz_to_dtype(tz): +def tz_to_dtype(tz: tzinfo | None, unit: str = "ns"): """ Return a datetime64[ns] dtype appropriate for the given timezone. Parameters ---------- tz : tzinfo or None + unit : str, default "ns" Returns ------- np.dtype or Datetime64TZDType """ if tz is None: - return DT64NS_DTYPE + return np.dtype(f"M8[{unit}]") else: - return DatetimeTZDtype(tz=tz) + return DatetimeTZDtype(tz=tz, unit=unit) def _field_accessor(name: str, field: str, docstring=None): @@ -800,7 +801,7 @@ def tz_convert(self, tz) -> DatetimeArray: ) # No conversion since timestamps are all UTC to begin with - dtype = tz_to_dtype(tz) + dtype = tz_to_dtype(tz, unit=self._unit) return self._simple_new(self._ndarray, dtype=dtype, freq=self.freq) @dtl.ravel_compat @@ -965,10 +966,14 @@ def tz_localize(self, tz, ambiguous="raise", nonexistent="raise") -> DatetimeArr # Convert to UTC new_dates = tzconversion.tz_localize_to_utc( - self.asi8, tz, ambiguous=ambiguous, nonexistent=nonexistent + self.asi8, + tz, + ambiguous=ambiguous, + nonexistent=nonexistent, + reso=self._reso, ) - new_dates = new_dates.view(DT64NS_DTYPE) - dtype = tz_to_dtype(tz) + new_dates = new_dates.view(f"M8[{self._unit}]") + dtype = tz_to_dtype(tz, unit=self._unit) freq = None if timezones.is_utc(tz) or (len(self) == 1 and not isna(new_dates[0])): diff --git a/pandas/tests/arrays/test_timedeltas.py b/pandas/tests/arrays/test_timedeltas.py index abc27469a5428..b3b79bd988ad8 100644 --- a/pandas/tests/arrays/test_timedeltas.py +++ b/pandas/tests/arrays/test_timedeltas.py @@ -92,6 +92,34 @@ def test_add_pdnat(self, tda): assert result._reso == tda._reso assert result.isna().all() + # TODO: 2022-07-11 this is the only test that gets to DTA.tz_convert + # or tz_localize with non-nano; implement tests specific to that. + def test_add_datetimelike_scalar(self, tda, tz_naive_fixture): + ts = pd.Timestamp("2016-01-01", tz=tz_naive_fixture) + + msg = "with mis-matched resolutions" + with pytest.raises(NotImplementedError, match=msg): + # mismatched reso -> check that we don't give an incorrect result + tda + ts + with pytest.raises(NotImplementedError, match=msg): + # mismatched reso -> check that we don't give an incorrect result + ts + tda + + ts = ts._as_unit(tda._unit) + + exp_values = tda._ndarray + ts.asm8 + expected = ( + DatetimeArray._simple_new(exp_values, dtype=exp_values.dtype) + .tz_localize("UTC") + .tz_convert(ts.tz) + ) + + result = tda + ts + tm.assert_extension_array_equal(result, expected) + + result = ts + tda + tm.assert_extension_array_equal(result, expected) + def test_mul_scalar(self, tda): other = 2 result = tda * other