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