diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index ac75e5ae5e2a0..a68d883f04380 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -2,7 +2,7 @@ Base and utility classes for tseries type pandas objects. """ import warnings - +import operator from datetime import datetime, timedelta from pandas import compat @@ -10,6 +10,12 @@ 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,13 +683,24 @@ 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 @@ -674,12 +708,6 @@ def __add__(self, other): 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,13 +737,26 @@ 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. @@ -723,14 +764,6 @@ def __sub__(self, other): .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__, diff --git a/pandas/core/indexes/datetimes.py b/pandas/core/indexes/datetimes.py index 17f92339e4205..36ea2bffb9531 100644 --- a/pandas/core/indexes/datetimes.py +++ b/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) diff --git a/pandas/core/indexes/period.py b/pandas/core/indexes/period.py index 4c14cbffcd813..f0567c9c963af 100644 --- a/pandas/core/indexes/period.py +++ b/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 diff --git a/pandas/core/indexes/timedeltas.py b/pandas/core/indexes/timedeltas.py index 3542a24290f89..219adfdb66c82 100644 --- a/pandas/core/indexes/timedeltas.py +++ b/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):