ENH: implement Index.__array_ufunc__ (pandas-dev#43904)

Author: jbrockmendel

doc/source/whatsnew/v1.4.0.rst

@@ -390,6 +390,7 @@ Datetimelike
 - Bug in :func:`to_datetime` with ``format`` and ``pandas.NA`` was raising ``ValueError`` (:issue:`42957`)
 - :func:`to_datetime` would silently swap ``MM/DD/YYYY`` and ``DD/MM/YYYY`` formats if the given ``dayfirst`` option could not be respected - now, a warning is raised in the case of delimited date strings (e.g. ``31-12-2012``) (:issue:`12585`)
 - Bug in :meth:`date_range` and :meth:`bdate_range` do not return right bound when ``start`` = ``end`` and set is closed on one side (:issue:`43394`)
+- Bug in inplace addition and subtraction of :class:`DatetimeIndex` or :class:`TimedeltaIndex` with :class:`DatetimeArray` or :class:`TimedeltaArray` (:issue:`43904`)
 -
 
 Timedelta

pandas/core/arraylike.py

@@ -357,7 +357,7 @@ def reconstruct(result):
         return result
 
     if "out" in kwargs:
-        result = _dispatch_ufunc_with_out(self, ufunc, method, *inputs, **kwargs)
+        result = dispatch_ufunc_with_out(self, ufunc, method, *inputs, **kwargs)
         return reconstruct(result)
 
     # We still get here with kwargs `axis` for e.g. np.maximum.accumulate
@@ -410,7 +410,7 @@ def _standardize_out_kwarg(**kwargs) -> dict:
     return kwargs
 
 
-def _dispatch_ufunc_with_out(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
+def dispatch_ufunc_with_out(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
     """
     If we have an `out` keyword, then call the ufunc without `out` and then
     set the result into the given `out`.

pandas/core/arrays/base.py

@@ -1379,6 +1379,11 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
         if result is not NotImplemented:
             return result
 
+        if "out" in kwargs:
+            return arraylike.dispatch_ufunc_with_out(
+                self, ufunc, method, *inputs, **kwargs
+            )
+
         return arraylike.default_array_ufunc(self, ufunc, method, *inputs, **kwargs)
 
 

pandas/core/arrays/datetimelike.py

@@ -1414,7 +1414,7 @@ def __iadd__(self, other):
 
         if not is_period_dtype(self.dtype):
             # restore freq, which is invalidated by setitem
-            self._freq = result._freq
+            self._freq = result.freq
         return self
 
     def __isub__(self, other):
@@ -1423,7 +1423,7 @@ def __isub__(self, other):
 
         if not is_period_dtype(self.dtype):
             # restore freq, which is invalidated by setitem
-            self._freq = result._freq
+            self._freq = result.freq
         return self
 
     # --------------------------------------------------------------

pandas/core/indexes/base.py

@@ -102,6 +102,7 @@
     PeriodDtype,
 )
 from pandas.core.dtypes.generic import (
+    ABCDataFrame,
     ABCDatetimeIndex,
     ABCMultiIndex,
     ABCPeriodIndex,
@@ -116,6 +117,7 @@
 )
 
 from pandas.core import (
+    arraylike,
     missing,
     ops,
 )
@@ -844,6 +846,24 @@ def __array__(self, dtype=None) -> np.ndarray:
         """
         return np.asarray(self._data, dtype=dtype)
 
+    def __array_ufunc__(self, ufunc: np.ufunc, method: str_t, *inputs, **kwargs):
+        if any(isinstance(other, (ABCSeries, ABCDataFrame)) for other in inputs):
+            return NotImplemented
+
+        result = arraylike.maybe_dispatch_ufunc_to_dunder_op(
+            self, ufunc, method, *inputs, **kwargs
+        )
+        if result is not NotImplemented:
+            return result
+
+        new_inputs = [x if x is not self else x._values for x in inputs]
+        result = getattr(ufunc, method)(*new_inputs, **kwargs)
+        if ufunc.nout == 2:
+            # i.e. np.divmod, np.modf, np.frexp
+            return tuple(self.__array_wrap__(x) for x in result)
+
+        return self.__array_wrap__(result)
+
     def __array_wrap__(self, result, context=None):
         """
         Gets called after a ufunc and other functions.

pandas/core/indexes/datetimelike.py

@@ -672,15 +672,6 @@ def insert(self, loc: int, item):
     # --------------------------------------------------------------------
     # NDArray-Like Methods
 
-    def __array_wrap__(self, result, context=None):
-        """
-        Gets called after a ufunc and other functions.
-        """
-        out = super().__array_wrap__(result, context=context)
-        if isinstance(out, DatetimeTimedeltaMixin) and self.freq is not None:
-            out = out._with_freq("infer")
-        return out
-
     @Appender(_index_shared_docs["take"] % _index_doc_kwargs)
     def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs):
         nv.validate_take((), kwargs)

pandas/tests/arithmetic/test_datetime64.py

@@ -2163,6 +2163,15 @@ def test_dti_isub_tdi(self, tz_naive_fixture):
         result -= tdi
         tm.assert_index_equal(result, expected)
 
+        # DTA.__isub__ GH#43904
+        dta = dti._data.copy()
+        dta -= tdi
+        tm.assert_datetime_array_equal(dta, expected._data)
+
+        out = dti._data.copy()
+        np.subtract(out, tdi, out=out)
+        tm.assert_datetime_array_equal(out, expected._data)
+
         msg = "cannot subtract .* from a TimedeltaArray"
         with pytest.raises(TypeError, match=msg):
             tdi -= dti
@@ -2172,10 +2181,14 @@ def test_dti_isub_tdi(self, tz_naive_fixture):
         result -= tdi.values
         tm.assert_index_equal(result, expected)
 
-        msg = "cannot subtract a datelike from a TimedeltaArray"
+        msg = "cannot subtract DatetimeArray from ndarray"
         with pytest.raises(TypeError, match=msg):
             tdi.values -= dti
 
+        msg = "cannot subtract a datelike from a TimedeltaArray"
+        with pytest.raises(TypeError, match=msg):
+            tdi._values -= dti
+
     # -------------------------------------------------------------
     # Binary Operations DatetimeIndex and datetime-like
     # TODO: A couple other tests belong in this section.  Move them in