De-duplicate add_offset_array methods (#19835)

Author: jbrockmendel

pandas/core/indexes/datetimelike.py

@@ -2,14 +2,20 @@
 Base and utility classes for tseries type pandas objects.
 """
 import warnings
-
+import operator
 from datetime import datetime, timedelta
 
 from pandas import compat
 from pandas.compat.numpy import function as nv
 from pandas.core.tools.timedeltas import to_timedelta
 
 import numpy as np
+
+from pandas._libs import lib, iNaT, NaT
+from pandas._libs.tslibs.period import Period
+from pandas._libs.tslibs.timedeltas import delta_to_nanoseconds
+from pandas._libs.tslibs.timestamps import round_ns
+
 from pandas.core.dtypes.common import (
     _ensure_int64,
     is_dtype_equal,
@@ -25,18 +31,15 @@
     is_integer_dtype,
     is_object_dtype,
     is_string_dtype,
+    is_period_dtype,
     is_timedelta64_dtype)
 from pandas.core.dtypes.generic import (
     ABCIndex, ABCSeries, ABCPeriodIndex, ABCIndexClass)
 from pandas.core.dtypes.missing import isna
 from pandas.core import common as com, algorithms, ops
 from pandas.core.algorithms import checked_add_with_arr
-from pandas.errors import NullFrequencyError
+from pandas.errors import NullFrequencyError, PerformanceWarning
 import pandas.io.formats.printing as printing
-from pandas._libs import lib, iNaT, NaT
-from pandas._libs.tslibs.period import Period
-from pandas._libs.tslibs.timedeltas import delta_to_nanoseconds
-from pandas._libs.tslibs.timestamps import round_ns
 
 from pandas.core.indexes.base import Index, _index_shared_docs
 from pandas.util._decorators import Appender, cache_readonly
@@ -637,13 +640,33 @@ def _sub_datelike(self, other):
     def _sub_period(self, other):
         return NotImplemented
 
-    def _add_offset_array(self, other):
-        # Array/Index of DateOffset objects
-        return NotImplemented
+    def _addsub_offset_array(self, other, op):
+        """
+        Add or subtract array-like of DateOffset objects
 
-    def _sub_offset_array(self, other):
-        # Array/Index of DateOffset objects
-        return NotImplemented
+        Parameters
+        ----------
+        other : Index, np.ndarray
+            object-dtype containing pd.DateOffset objects
+        op : {operator.add, operator.sub}
+
+        Returns
+        -------
+        result : same class as self
+        """
+        assert op in [operator.add, operator.sub]
+        if len(other) == 1:
+            return op(self, other[0])
+
+        warnings.warn("Adding/subtracting array of DateOffsets to "
+                      "{cls} not vectorized"
+                      .format(cls=type(self).__name__), PerformanceWarning)
+
+        res_values = op(self.astype('O').values, np.array(other))
+        kwargs = {}
+        if not is_period_dtype(self):
+            kwargs['freq'] = 'infer'
+        return self._constructor(res_values, **kwargs)
 
     @classmethod
     def _add_datetimelike_methods(cls):
@@ -660,26 +683,31 @@ def __add__(self, other):
             other = lib.item_from_zerodim(other)
             if isinstance(other, ABCSeries):
                 return NotImplemented
-            elif is_timedelta64_dtype(other):
+
+            # scalar others
+            elif isinstance(other, (DateOffset, timedelta, np.timedelta64)):
                 result = self._add_delta(other)
-            elif isinstance(other, (DateOffset, timedelta)):
+            elif isinstance(other, (datetime, np.datetime64)):
+                result = self._add_datelike(other)
+            elif is_integer(other):
+                # This check must come after the check for np.timedelta64
+                # as is_integer returns True for these
+                result = self.shift(other)
+
+            # array-like others
+            elif is_timedelta64_dtype(other):
+                # TimedeltaIndex, ndarray[timedelta64]
                 result = self._add_delta(other)
             elif is_offsetlike(other):
                 # Array/Index of DateOffset objects
-                result = self._add_offset_array(other)
+                result = self._addsub_offset_array(other, operator.add)
             elif isinstance(self, TimedeltaIndex) and isinstance(other, Index):
                 if hasattr(other, '_add_delta'):
                     # i.e. DatetimeIndex, TimedeltaIndex, or PeriodIndex
                     result = other._add_delta(self)
                 else:
                     raise TypeError("cannot add TimedeltaIndex and {typ}"
                                     .format(typ=type(other)))
-            elif is_integer(other):
-                # This check must come after the check for timedelta64_dtype
-                # or else it will incorrectly catch np.timedelta64 objects
-                result = self.shift(other)
-            elif isinstance(other, (datetime, np.datetime64)):
-                result = self._add_datelike(other)
             elif isinstance(other, Index):
                 result = self._add_datelike(other)
             elif is_integer_dtype(other) and self.freq is None:
@@ -709,28 +737,33 @@ def __sub__(self, other):
             other = lib.item_from_zerodim(other)
             if isinstance(other, ABCSeries):
                 return NotImplemented
-            elif is_timedelta64_dtype(other):
+
+            # scalar others
+            elif isinstance(other, (DateOffset, timedelta, np.timedelta64)):
                 result = self._add_delta(-other)
-            elif isinstance(other, (DateOffset, timedelta)):
+            elif isinstance(other, (datetime, np.datetime64)):
+                result = self._sub_datelike(other)
+            elif is_integer(other):
+                # This check must come after the check for np.timedelta64
+                # as is_integer returns True for these
+                result = self.shift(-other)
+            elif isinstance(other, Period):
+                result = self._sub_period(other)
+
+            # array-like others
+            elif is_timedelta64_dtype(other):
+                # TimedeltaIndex, ndarray[timedelta64]
                 result = self._add_delta(-other)
             elif is_offsetlike(other):
                 # Array/Index of DateOffset objects
-                result = self._sub_offset_array(other)
+                result = self._addsub_offset_array(other, operator.sub)
             elif isinstance(self, TimedeltaIndex) and isinstance(other, Index):
                 # We checked above for timedelta64_dtype(other) so this
                 # must be invalid.
                 raise TypeError("cannot subtract TimedeltaIndex and {typ}"
                                 .format(typ=type(other).__name__))
             elif isinstance(other, DatetimeIndex):
                 result = self._sub_datelike(other)
-            elif is_integer(other):
-                # This check must come after the check for timedelta64_dtype
-                # or else it will incorrectly catch np.timedelta64 objects
-                result = self.shift(-other)
-            elif isinstance(other, (datetime, np.datetime64)):
-                result = self._sub_datelike(other)
-            elif isinstance(other, Period):
-                result = self._sub_period(other)
             elif isinstance(other, Index):
                 raise TypeError("cannot subtract {typ1} and {typ2}"
                                 .format(typ1=type(self).__name__,

pandas/core/indexes/datetimes.py

@@ -964,29 +964,6 @@ def _add_offset(self, offset):
                           "or DatetimeIndex", PerformanceWarning)
             return self.astype('O') + offset
 
-    def _add_offset_array(self, other):
-        # Array/Index of DateOffset objects
-        if len(other) == 1:
-            return self + other[0]
-        else:
-            warnings.warn("Adding/subtracting array of DateOffsets to "
-                          "{} not vectorized".format(type(self)),
-                          PerformanceWarning)
-            return self.astype('O') + np.array(other)
-            # TODO: pass freq='infer' like we do in _sub_offset_array?
-            # TODO: This works for __add__ but loses dtype in __sub__
-
-    def _sub_offset_array(self, other):
-        # Array/Index of DateOffset objects
-        if len(other) == 1:
-            return self - other[0]
-        else:
-            warnings.warn("Adding/subtracting array of DateOffsets to "
-                          "{} not vectorized".format(type(self)),
-                          PerformanceWarning)
-            res_values = self.astype('O').values - np.array(other)
-            return self.__class__(res_values, freq='infer')
-
     def _format_native_types(self, na_rep='NaT', date_format=None, **kwargs):
         from pandas.io.formats.format import _get_format_datetime64_from_values
         format = _get_format_datetime64_from_values(self, date_format)

pandas/core/indexes/period.py

@@ -44,7 +44,6 @@
 from pandas.util._decorators import (Appender, Substitution, cache_readonly,
                                      deprecate_kwarg)
 from pandas.compat import zip, u
-from pandas.errors import PerformanceWarning
 
 import pandas.core.indexes.base as ibase
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
@@ -745,28 +744,6 @@ def _sub_period(self, other):
         # result must be Int64Index or Float64Index
         return Index(new_data)
 
-    def _add_offset_array(self, other):
-        # Array/Index of DateOffset objects
-        if len(other) == 1:
-            return self + other[0]
-        else:
-            warnings.warn("Adding/subtracting array of DateOffsets to "
-                          "{cls} not vectorized"
-                          .format(cls=type(self).__name__), PerformanceWarning)
-            res_values = self.astype('O').values + np.array(other)
-            return self.__class__(res_values)
-
-    def _sub_offset_array(self, other):
-        # Array/Index of DateOffset objects
-        if len(other) == 1:
-            return self - other[0]
-        else:
-            warnings.warn("Adding/subtracting array of DateOffsets to "
-                          "{cls} not vectorized"
-                          .format(cls=type(self).__name__), PerformanceWarning)
-            res_values = self.astype('O').values - np.array(other)
-            return self.__class__(res_values)
-
     def shift(self, n):
         """
         Specialized shift which produces an PeriodIndex

pandas/core/indexes/timedeltas.py

@@ -1,7 +1,6 @@
 """ implement the TimedeltaIndex """
 
 from datetime import timedelta
-import warnings
 
 import numpy as np
 from pandas.core.dtypes.common import (
@@ -433,43 +432,16 @@ def _sub_datelike(self, other):
         else:
             raise TypeError("cannot subtract a datelike from a TimedeltaIndex")
 
-    def _add_offset_array(self, other):
-        # Array/Index of DateOffset objects
+    def _addsub_offset_array(self, other, op):
+        # Add or subtract Array-like of DateOffset objects
         try:
             # TimedeltaIndex can only operate with a subset of DateOffset
             # subclasses.  Incompatible classes will raise AttributeError,
             # which we re-raise as TypeError
-            if len(other) == 1:
-                return self + other[0]
-            else:
-                from pandas.errors import PerformanceWarning
-                warnings.warn("Adding/subtracting array of DateOffsets to "
-                              "{} not vectorized".format(type(self)),
-                              PerformanceWarning)
-                return self.astype('O') + np.array(other)
-                # TODO: pass freq='infer' like we do in _sub_offset_array?
-                # TODO: This works for __add__ but loses dtype in __sub__
-        except AttributeError:
-            raise TypeError("Cannot add non-tick DateOffset to TimedeltaIndex")
-
-    def _sub_offset_array(self, other):
-        # Array/Index of DateOffset objects
-        try:
-            # TimedeltaIndex can only operate with a subset of DateOffset
-            # subclasses.  Incompatible classes will raise AttributeError,
-            # which we re-raise as TypeError
-            if len(other) == 1:
-                return self - other[0]
-            else:
-                from pandas.errors import PerformanceWarning
-                warnings.warn("Adding/subtracting array of DateOffsets to "
-                              "{} not vectorized".format(type(self)),
-                              PerformanceWarning)
-                res_values = self.astype('O').values - np.array(other)
-                return self.__class__(res_values, freq='infer')
+            return DatetimeIndexOpsMixin._addsub_offset_array(self, other, op)
         except AttributeError:
-            raise TypeError("Cannot subtrack non-tick DateOffset from"
-                            " TimedeltaIndex")
+            raise TypeError("Cannot add/subtract non-tick DateOffset to {cls}"
+                            .format(cls=type(self).__name__))
 
     def _format_native_types(self, na_rep=u('NaT'),
                              date_format=None, **kwargs):