REF: move range-generation functions to EA mixin classes (#22016)

Author: jbrockmendel

pandas/core/arrays/datetimelike.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+from datetime import datetime, timedelta
 import operator
 import warnings
 
@@ -8,7 +9,7 @@
 from pandas._libs.tslibs import timezones
 from pandas._libs.tslibs.timedeltas import delta_to_nanoseconds, Timedelta
 from pandas._libs.tslibs.period import (
-    DIFFERENT_FREQ_INDEX, IncompatibleFrequency)
+    Period, DIFFERENT_FREQ_INDEX, IncompatibleFrequency)
 
 from pandas.errors import NullFrequencyError, PerformanceWarning
 from pandas import compat
@@ -19,6 +20,13 @@
 from pandas.core.dtypes.common import (
     needs_i8_conversion,
     is_list_like,
+    is_offsetlike,
+    is_extension_array_dtype,
+    is_datetime64_dtype,
+    is_datetime64_any_dtype,
+    is_datetime64tz_dtype,
+    is_float_dtype,
+    is_integer_dtype,
     is_bool_dtype,
     is_period_dtype,
     is_timedelta64_dtype,
@@ -100,7 +108,7 @@ class DatetimeLikeArrayMixin(ExtensionOpsMixin, AttributesMixin):
         _freq
 
     and that the inheriting class has methods:
-        _validate_frequency
+        _generate_range
     """
 
     @property
@@ -132,6 +140,14 @@ def asi8(self):
     # ------------------------------------------------------------------
     # Array-like Methods
 
+    @property
+    def shape(self):
+        return (len(self),)
+
+    @property
+    def size(self):
+        return np.prod(self.shape)
+
     def __len__(self):
         return len(self._data)
 
@@ -296,6 +312,34 @@ def resolution(self):
         """
         return frequencies.Resolution.get_str(self._resolution)
 
+    @classmethod
+    def _validate_frequency(cls, index, freq, **kwargs):
+        """
+        Validate that a frequency is compatible with the values of a given
+        Datetime Array/Index or Timedelta Array/Index
+
+        Parameters
+        ----------
+        index : DatetimeIndex or TimedeltaIndex
+            The index on which to determine if the given frequency is valid
+        freq : DateOffset
+            The frequency to validate
+        """
+        if is_period_dtype(cls):
+            # Frequency validation is not meaningful for Period Array/Index
+            return None
+
+        inferred = index.inferred_freq
+        if index.size == 0 or inferred == freq.freqstr:
+            return None
+
+        on_freq = cls._generate_range(start=index[0], end=None,
+                                      periods=len(index), freq=freq, **kwargs)
+        if not np.array_equal(index.asi8, on_freq.asi8):
+            raise ValueError('Inferred frequency {infer} from passed values '
+                             'does not conform to passed frequency {passed}'
+                             .format(infer=inferred, passed=freq.freqstr))
+
     # ------------------------------------------------------------------
     # Arithmetic Methods
 
@@ -477,6 +521,188 @@ def _addsub_offset_array(self, other, op):
             kwargs['freq'] = 'infer'
         return type(self)(res_values, **kwargs)
 
+    def shift(self, n, freq=None):
+        """
+        Specialized shift which produces a Datetime/Timedelta Array/Index
+
+        Parameters
+        ----------
+        n : int
+            Periods to shift by
+        freq : DateOffset or timedelta-like, optional
+
+        Returns
+        -------
+        shifted : same type as self
+        """
+        if freq is not None and freq != self.freq:
+            if isinstance(freq, compat.string_types):
+                freq = frequencies.to_offset(freq)
+            offset = n * freq
+            result = self + offset
+
+            if hasattr(self, 'tz'):
+                result._tz = self.tz
+
+            return result
+
+        if n == 0:
+            # immutable so OK
+            return self
+
+        if self.freq is None:
+            raise NullFrequencyError("Cannot shift with no freq")
+
+        start = self[0] + n * self.freq
+        end = self[-1] + n * self.freq
+        attribs = self._get_attributes_dict()
+        return self._generate_range(start=start, end=end, periods=None,
+                                    **attribs)
+
+    @classmethod
+    def _add_datetimelike_methods(cls):
+        """
+        add in the datetimelike methods (as we may have to override the
+        superclass)
+        """
+
+        def __add__(self, other):
+            other = lib.item_from_zerodim(other)
+            if isinstance(other, (ABCSeries, ABCDataFrame)):
+                return NotImplemented
+
+            # scalar others
+            elif other is NaT:
+                result = self._add_nat()
+            elif isinstance(other, (Tick, timedelta, np.timedelta64)):
+                result = self._add_delta(other)
+            elif isinstance(other, DateOffset):
+                # specifically _not_ a Tick
+                result = self._add_offset(other)
+            elif isinstance(other, (datetime, np.datetime64)):
+                result = self._add_datelike(other)
+            elif lib.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._addsub_offset_array(other, operator.add)
+            elif is_datetime64_dtype(other) or is_datetime64tz_dtype(other):
+                # DatetimeIndex, ndarray[datetime64]
+                return self._add_datelike(other)
+            elif is_integer_dtype(other):
+                result = self._addsub_int_array(other, operator.add)
+            elif is_float_dtype(other) or is_period_dtype(other):
+                # Explicitly catch invalid dtypes
+                raise TypeError("cannot add {dtype}-dtype to {cls}"
+                                .format(dtype=other.dtype,
+                                        cls=type(self).__name__))
+            elif is_extension_array_dtype(other):
+                # Categorical op will raise; defer explicitly
+                return NotImplemented
+            else:  # pragma: no cover
+                return NotImplemented
+
+            return result
+
+        cls.__add__ = __add__
+
+        def __radd__(self, other):
+            # alias for __add__
+            return self.__add__(other)
+        cls.__radd__ = __radd__
+
+        def __sub__(self, other):
+            other = lib.item_from_zerodim(other)
+            if isinstance(other, (ABCSeries, ABCDataFrame)):
+                return NotImplemented
+
+            # scalar others
+            elif other is NaT:
+                result = self._sub_nat()
+            elif isinstance(other, (Tick, timedelta, np.timedelta64)):
+                result = self._add_delta(-other)
+            elif isinstance(other, DateOffset):
+                # specifically _not_ a Tick
+                result = self._add_offset(-other)
+            elif isinstance(other, (datetime, np.datetime64)):
+                result = self._sub_datelike(other)
+            elif lib.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._addsub_offset_array(other, operator.sub)
+            elif is_datetime64_dtype(other) or is_datetime64tz_dtype(other):
+                # DatetimeIndex, ndarray[datetime64]
+                result = self._sub_datelike(other)
+            elif is_period_dtype(other):
+                # PeriodIndex
+                result = self._sub_period_array(other)
+            elif is_integer_dtype(other):
+                result = self._addsub_int_array(other, operator.sub)
+            elif isinstance(other, ABCIndexClass):
+                raise TypeError("cannot subtract {cls} and {typ}"
+                                .format(cls=type(self).__name__,
+                                        typ=type(other).__name__))
+            elif is_float_dtype(other):
+                # Explicitly catch invalid dtypes
+                raise TypeError("cannot subtract {dtype}-dtype from {cls}"
+                                .format(dtype=other.dtype,
+                                        cls=type(self).__name__))
+            elif is_extension_array_dtype(other):
+                # Categorical op will raise; defer explicitly
+                return NotImplemented
+            else:  # pragma: no cover
+                return NotImplemented
+
+            return result
+
+        cls.__sub__ = __sub__
+
+        def __rsub__(self, other):
+            if is_datetime64_dtype(other) and is_timedelta64_dtype(self):
+                # ndarray[datetime64] cannot be subtracted from self, so
+                # we need to wrap in DatetimeArray/Index and flip the operation
+                if not isinstance(other, DatetimeLikeArrayMixin):
+                    # Avoid down-casting DatetimeIndex
+                    from pandas.core.arrays import DatetimeArrayMixin
+                    other = DatetimeArrayMixin(other)
+                return other - self
+            elif (is_datetime64_any_dtype(self) and hasattr(other, 'dtype') and
+                  not is_datetime64_any_dtype(other)):
+                # GH#19959 datetime - datetime is well-defined as timedelta,
+                # but any other type - datetime is not well-defined.
+                raise TypeError("cannot subtract {cls} from {typ}"
+                                .format(cls=type(self).__name__,
+                                        typ=type(other).__name__))
+            return -(self - other)
+        cls.__rsub__ = __rsub__
+
+        def __iadd__(self, other):
+            # alias for __add__
+            return self.__add__(other)
+        cls.__iadd__ = __iadd__
+
+        def __isub__(self, other):
+            # alias for __sub__
+            return self.__sub__(other)
+        cls.__isub__ = __isub__
+
     # --------------------------------------------------------------
     # Comparison Methods