ENH: TDA+datetime_scalar support non-nano (#47675)

Author: jbrockmendel

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:

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])):

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