Move most remaining arith helpers (#21815)

Author: jbrockmendel

pandas/core/arrays/datetimelike.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import operator
+import warnings
 
 import numpy as np
 
@@ -8,12 +9,16 @@
 from pandas._libs.tslibs.period import (
     DIFFERENT_FREQ_INDEX, IncompatibleFrequency)
 
-from pandas.errors import NullFrequencyError
+from pandas.errors import NullFrequencyError, PerformanceWarning
 
 from pandas.tseries import frequencies
 from pandas.tseries.offsets import Tick
 
-from pandas.core.dtypes.common import is_period_dtype, is_timedelta64_dtype
+from pandas.core.dtypes.common import (
+    is_period_dtype,
+    is_timedelta64_dtype,
+    is_object_dtype)
+
 import pandas.core.common as com
 from pandas.core.algorithms import checked_add_with_arr
 
@@ -108,38 +113,43 @@ def __getitem__(self, key):
         if is_int:
             val = getitem(key)
             return self._box_func(val)
+
+        if com.is_bool_indexer(key):
+            key = np.asarray(key)
+            if key.all():
+                key = slice(0, None, None)
+            else:
+                key = lib.maybe_booleans_to_slice(key.view(np.uint8))
+
+        attribs = self._get_attributes_dict()
+
+        is_period = is_period_dtype(self)
+        if is_period:
+            freq = self.freq
         else:
-            if com.is_bool_indexer(key):
-                key = np.asarray(key)
-                if key.all():
-                    key = slice(0, None, None)
+            freq = None
+            if isinstance(key, slice):
+                if self.freq is not None and key.step is not None:
+                    freq = key.step * self.freq
                 else:
-                    key = lib.maybe_booleans_to_slice(key.view(np.uint8))
+                    freq = self.freq
 
-            attribs = self._get_attributes_dict()
+        attribs['freq'] = freq
 
-            is_period = is_period_dtype(self)
+        result = getitem(key)
+        if result.ndim > 1:
+            # To support MPL which performs slicing with 2 dim
+            # even though it only has 1 dim by definition
             if is_period:
-                freq = self.freq
-            else:
-                freq = None
-                if isinstance(key, slice):
-                    if self.freq is not None and key.step is not None:
-                        freq = key.step * self.freq
-                    else:
-                        freq = self.freq
-
-            attribs['freq'] = freq
+                return self._simple_new(result, **attribs)
+            return result
 
-            result = getitem(key)
-            if result.ndim > 1:
-                # To support MPL which performs slicing with 2 dim
-                # even though it only has 1 dim by definition
-                if is_period:
-                    return self._simple_new(result, **attribs)
-                return result
+        return self._simple_new(result, **attribs)
 
-            return self._simple_new(result, **attribs)
+    def astype(self, dtype, copy=True):
+        if is_object_dtype(dtype):
+            return self._box_values(self.asi8)
+        return super(DatetimeLikeArrayMixin, self).astype(dtype, copy)
 
     # ------------------------------------------------------------------
     # Null Handling
@@ -397,3 +407,31 @@ def _addsub_int_array(self, other, op):
         # to _addsub_offset_array
         assert not is_timedelta64_dtype(self)
         return op(self, np.array(other) * self.freq)
+
+    def _addsub_offset_array(self, other, op):
+        """
+        Add or subtract array-like of DateOffset objects
+
+        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 type(self)(res_values, **kwargs)

pandas/core/arrays/datetimes.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+from datetime import timedelta
 import warnings
 
 import numpy as np
@@ -11,15 +12,18 @@
     resolution as libresolution)
 
 from pandas.util._decorators import cache_readonly
+from pandas.errors import PerformanceWarning
 
 from pandas.core.dtypes.common import (
     _NS_DTYPE,
     is_datetime64tz_dtype,
     is_datetime64_dtype,
+    is_timedelta64_dtype,
     _ensure_int64)
 from pandas.core.dtypes.dtypes import DatetimeTZDtype
 
 from pandas.tseries.frequencies import to_offset, DateOffset
+from pandas.tseries.offsets import Tick
 
 from .datetimelike import DatetimeLikeArrayMixin
 
@@ -104,6 +108,10 @@ def _simple_new(cls, values, freq=None, tz=None, **kwargs):
         return result
 
     def __new__(cls, values, freq=None, tz=None):
+        if tz is None and hasattr(values, 'tz'):
+            # e.g. DatetimeIndex
+            tz = values.tz
+
         if (freq is not None and not isinstance(freq, DateOffset) and
                 freq != 'infer'):
             freq = to_offset(freq)
@@ -131,6 +139,17 @@ def dtype(self):
             return _NS_DTYPE
         return DatetimeTZDtype('ns', self.tz)
 
+    @property
+    def tz(self):
+        # GH 18595
+        return self._tz
+
+    @tz.setter
+    def tz(self, value):
+        # GH 3746: Prevent localizing or converting the index by setting tz
+        raise AttributeError("Cannot directly set timezone. Use tz_localize() "
+                             "or tz_convert() as appropriate")
+
     @property
     def tzinfo(self):
         """
@@ -244,6 +263,60 @@ def _sub_datelike_dti(self, other):
             new_values[mask] = iNaT
         return new_values.view('timedelta64[ns]')
 
+    def _add_offset(self, offset):
+        assert not isinstance(offset, Tick)
+        try:
+            if self.tz is not None:
+                values = self.tz_localize(None)
+            else:
+                values = self
+            result = offset.apply_index(values)
+            if self.tz is not None:
+                result = result.tz_localize(self.tz)
+
+        except NotImplementedError:
+            warnings.warn("Non-vectorized DateOffset being applied to Series "
+                          "or DatetimeIndex", PerformanceWarning)
+            result = self.astype('O') + offset
+
+        return type(self)(result, freq='infer')
+
+    def _add_delta(self, delta):
+        """
+        Add a timedelta-like, DateOffset, or TimedeltaIndex-like object
+        to self.
+
+        Parameters
+        ----------
+        delta : {timedelta, np.timedelta64, DateOffset,
+                 TimedelaIndex, ndarray[timedelta64]}
+
+        Returns
+        -------
+        result : same type as self
+
+        Notes
+        -----
+        The result's name is set outside of _add_delta by the calling
+        method (__add__ or __sub__)
+        """
+        from pandas.core.arrays.timedelta import TimedeltaArrayMixin
+
+        if isinstance(delta, (Tick, timedelta, np.timedelta64)):
+            new_values = self._add_delta_td(delta)
+        elif is_timedelta64_dtype(delta):
+            if not isinstance(delta, TimedeltaArrayMixin):
+                delta = TimedeltaArrayMixin(delta)
+            new_values = self._add_delta_tdi(delta)
+        else:
+            new_values = self.astype('O') + delta
+
+        tz = 'UTC' if self.tz is not None else None
+        result = type(self)(new_values, tz=tz, freq='infer')
+        if self.tz is not None and self.tz is not utc:
+            result = result.tz_convert(self.tz)
+        return result
+
     # -----------------------------------------------------------------
     # Timezone Conversion and Localization Methods
 

pandas/core/arrays/period.py

@@ -8,14 +8,16 @@
 from pandas._libs.tslib import NaT, iNaT
 from pandas._libs.tslibs.period import (
     Period, IncompatibleFrequency, DIFFERENT_FREQ_INDEX,
-    get_period_field_arr)
+    get_period_field_arr, period_asfreq_arr)
+from pandas._libs.tslibs import period as libperiod
 from pandas._libs.tslibs.timedeltas import delta_to_nanoseconds
 from pandas._libs.tslibs.fields import isleapyear_arr
 
 from pandas import compat
 from pandas.util._decorators import cache_readonly
 
-from pandas.core.dtypes.common import is_integer_dtype, is_float_dtype
+from pandas.core.dtypes.common import (
+    is_integer_dtype, is_float_dtype, is_period_dtype)
 from pandas.core.dtypes.dtypes import PeriodDtype
 
 from pandas.tseries import frequencies
@@ -113,12 +115,23 @@ def freq(self, value):
 
     _attributes = ["freq"]
 
+    def __new__(cls, values, freq=None, **kwargs):
+        if is_period_dtype(values):
+            # PeriodArray, PeriodIndex
+            if freq is not None and values.freq != freq:
+                raise IncompatibleFrequency(freq, values.freq)
+            freq = values.freq
+            values = values.asi8
+
+        return cls._simple_new(values, freq, **kwargs)
+
     @classmethod
     def _simple_new(cls, values, freq=None, **kwargs):
         """
         Values can be any type that can be coerced to Periods.
         Ordinals in an ndarray are fastpath-ed to `_from_ordinals`
         """
+
         if not is_integer_dtype(values):
             values = np.array(values, copy=False)
             if len(values) > 0 and is_float_dtype(values):
@@ -128,8 +141,6 @@ def _simple_new(cls, values, freq=None, **kwargs):
 
         return cls._from_ordinals(values, freq)
 
-    __new__ = _simple_new  # For now...
-
     @classmethod
     def _from_ordinals(cls, values, freq=None):
         """
@@ -173,6 +184,65 @@ def is_leap_year(self):
         """ Logical indicating if the date belongs to a leap year """
         return isleapyear_arr(np.asarray(self.year))
 
+    def asfreq(self, freq=None, how='E'):
+        """
+        Convert the Period Array/Index to the specified frequency `freq`.
+
+        Parameters
+        ----------
+        freq : str
+            a frequency
+        how : str {'E', 'S'}
+            'E', 'END', or 'FINISH' for end,
+            'S', 'START', or 'BEGIN' for start.
+            Whether the elements should be aligned to the end
+            or start within pa period. January 31st ('END') vs.
+            January 1st ('START') for example.
+
+        Returns
+        -------
+        new : Period Array/Index with the new frequency
+
+        Examples
+        --------
+        >>> pidx = pd.period_range('2010-01-01', '2015-01-01', freq='A')
+        >>> pidx
+        <class 'pandas.core.indexes.period.PeriodIndex'>
+        [2010, ..., 2015]
+        Length: 6, Freq: A-DEC
+
+        >>> pidx.asfreq('M')
+        <class 'pandas.core.indexes.period.PeriodIndex'>
+        [2010-12, ..., 2015-12]
+        Length: 6, Freq: M
+
+        >>> pidx.asfreq('M', how='S')
+        <class 'pandas.core.indexes.period.PeriodIndex'>
+        [2010-01, ..., 2015-01]
+        Length: 6, Freq: M
+        """
+        how = libperiod._validate_end_alias(how)
+
+        freq = Period._maybe_convert_freq(freq)
+
+        base1, mult1 = frequencies.get_freq_code(self.freq)
+        base2, mult2 = frequencies.get_freq_code(freq)
+
+        asi8 = self.asi8
+        # mult1 can't be negative or 0
+        end = how == 'E'
+        if end:
+            ordinal = asi8 + mult1 - 1
+        else:
+            ordinal = asi8
+
+        new_data = period_asfreq_arr(ordinal, base1, base2, end)
+
+        if self.hasnans:
+            new_data[self._isnan] = iNaT
+
+        return self._simple_new(new_data, self.name, freq=freq)
+
     # ------------------------------------------------------------------
     # Arithmetic Methods