ENH: DTA/TDA add datetimelike scalar with mismatched reso (pandas-dev#48669)

* ENH: DTA/TDA add datetimelike scalar with mismatched reso

* mypy fixup

Author: jbrockmendel

pandas/_libs/tslibs/__init__.py

@@ -31,13 +31,15 @@
     "periods_per_day",
     "periods_per_second",
     "is_supported_unit",
+    "npy_unit_to_abbrev",
 ]
 
 from pandas._libs.tslibs import dtypes
 from pandas._libs.tslibs.conversion import localize_pydatetime
 from pandas._libs.tslibs.dtypes import (
     Resolution,
     is_supported_unit,
+    npy_unit_to_abbrev,
     periods_per_day,
     periods_per_second,
 )

pandas/core/arrays/datetimelike.py

@@ -37,11 +37,13 @@
     Resolution,
     Tick,
     Timestamp,
+    astype_overflowsafe,
     delta_to_nanoseconds,
     get_unit_from_dtype,
     iNaT,
     ints_to_pydatetime,
     ints_to_pytimedelta,
+    npy_unit_to_abbrev,
     to_offset,
 )
 from pandas._libs.tslibs.fields import (
@@ -1130,10 +1132,13 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray:
             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."
-            )
+            # Just as with Timestamp/Timedelta, we cast to the lower resolution
+            #  so long as doing so is lossless.
+            if self._reso < other._reso:
+                other = other._as_unit(self._unit, round_ok=False)
+            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)
@@ -1289,10 +1294,12 @@ def _add_timedelta_arraylike(
         self = cast("DatetimeArray | TimedeltaArray", self)
 
         if self._reso != other._reso:
-            raise NotImplementedError(
-                f"Addition of {type(self).__name__} with TimedeltaArray with "
-                "mis-matched resolutions is not yet supported."
-            )
+            # Just as with Timestamp/Timedelta, we cast to the lower resolution
+            #  so long as doing so is lossless.
+            if self._reso < other._reso:
+                other = other._as_unit(self._unit)
+            else:
+                self = self._as_unit(other._unit)
 
         self_i8 = self.asi8
         other_i8 = other.asi8
@@ -2028,6 +2035,22 @@ def _unit(self) -> str:
         # "ExtensionDtype"; expected "Union[DatetimeTZDtype, dtype[Any]]"
         return dtype_to_unit(self.dtype)  # type: ignore[arg-type]
 
+    def _as_unit(self: TimelikeOpsT, unit: str) -> TimelikeOpsT:
+        dtype = np.dtype(f"{self.dtype.kind}8[{unit}]")
+        new_values = astype_overflowsafe(self._ndarray, dtype, round_ok=False)
+
+        if isinstance(self.dtype, np.dtype):
+            new_dtype = new_values.dtype
+        else:
+            tz = cast("DatetimeArray", self).tz
+            new_dtype = DatetimeTZDtype(tz=tz, unit=unit)
+
+        # error: Unexpected keyword argument "freq" for "_simple_new" of
+        # "NDArrayBacked"  [call-arg]
+        return type(self)._simple_new(
+            new_values, dtype=new_dtype, freq=self.freq  # type: ignore[call-arg]
+        )
+
     # --------------------------------------------------------------
 
     def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):

pandas/tests/arrays/test_timedeltas.py

@@ -104,11 +104,18 @@ def test_add_pdnat(self, tda):
     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):
+        expected = tda + ts._as_unit(tda._unit)
+        res = tda + ts
+        tm.assert_extension_array_equal(res, expected)
+        res = ts + tda
+        tm.assert_extension_array_equal(res, expected)
+
+        ts += Timedelta(1)  # so we can't cast losslessly
+        msg = "Cannot losslessly convert units"
+        with pytest.raises(ValueError, match=msg):
             # mismatched reso -> check that we don't give an incorrect result
             tda + ts
-        with pytest.raises(NotImplementedError, match=msg):
+        with pytest.raises(ValueError, match=msg):
             # mismatched reso -> check that we don't give an incorrect result
             ts + tda
 
@@ -179,13 +186,28 @@ def test_add_timedeltaarraylike(self, tda):
         tda_nano = TimedeltaArray(tda._ndarray.astype("m8[ns]"))
 
         msg = "mis-matched resolutions is not yet supported"
-        with pytest.raises(NotImplementedError, match=msg):
+        expected = tda * 2
+        res = tda_nano + tda
+        tm.assert_extension_array_equal(res, expected)
+        res = tda + tda_nano
+        tm.assert_extension_array_equal(res, expected)
+
+        expected = tda * 0
+        res = tda - tda_nano
+        tm.assert_extension_array_equal(res, expected)
+
+        res = tda_nano - tda
+        tm.assert_extension_array_equal(res, expected)
+
+        tda_nano[:] = np.timedelta64(1, "ns")  # can't round losslessly
+        msg = "Cannot losslessly cast '-?1 ns' to"
+        with pytest.raises(ValueError, match=msg):
             tda_nano + tda
-        with pytest.raises(NotImplementedError, match=msg):
+        with pytest.raises(ValueError, match=msg):
             tda + tda_nano
-        with pytest.raises(NotImplementedError, match=msg):
+        with pytest.raises(ValueError, match=msg):
             tda - tda_nano
-        with pytest.raises(NotImplementedError, match=msg):
+        with pytest.raises(ValueError, match=msg):
             tda_nano - tda
 
         result = tda_nano + tda_nano

pandas/tests/tslibs/test_api.py

@@ -56,6 +56,7 @@ def test_namespace():
         "periods_per_day",
         "periods_per_second",
         "is_supported_unit",
+        "npy_unit_to_abbrev",
     ]
 
     expected = set(submodules + api)