diff --git a/pandas/core/arrays/timedeltas.py b/pandas/core/arrays/timedeltas.py index 6a7ce7033efa0..c0cfa996810bc 100644 --- a/pandas/core/arrays/timedeltas.py +++ b/pandas/core/arrays/timedeltas.py @@ -2,7 +2,6 @@ from __future__ import division from datetime import timedelta -import operator import warnings import numpy as np @@ -17,13 +16,12 @@ from pandas.core.dtypes.common import ( _TD_DTYPE, ensure_int64, is_datetime64_dtype, is_float_dtype, - is_integer_dtype, is_list_like, is_object_dtype, is_string_dtype, - is_timedelta64_dtype) + is_integer_dtype, is_list_like, is_object_dtype, is_scalar, + is_string_dtype, is_timedelta64_dtype) from pandas.core.dtypes.generic import ( ABCDataFrame, ABCIndexClass, ABCSeries, ABCTimedeltaIndex) from pandas.core.dtypes.missing import isna -from pandas.core import ops from pandas.core.algorithms import checked_add_with_arr, unique1d import pandas.core.common as com @@ -106,29 +104,6 @@ def wrapper(self, other): return compat.set_function_name(wrapper, opname, cls) -def _wrap_tdi_op(op): - """ - Instead of re-implementing multiplication/division etc operations - in the Array class, for now we dispatch to the TimedeltaIndex - implementations. - """ - # TODO: implement directly here and wrap in TimedeltaIndex, instead of - # the other way around - def method(self, other): - if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)): - return NotImplemented - - from pandas import TimedeltaIndex - obj = TimedeltaIndex(self) - result = op(obj, other) - if is_timedelta64_dtype(result): - return type(self)(result) - return np.array(result) - - method.__name__ = '__{name}__'.format(name=op.__name__) - return method - - class TimedeltaArrayMixin(dtl.DatetimeLikeArrayMixin, dtl.TimelikeOps): _typ = "timedeltaarray" __array_priority__ = 1000 @@ -332,37 +307,41 @@ def _addsub_offset_array(self, other, op): raise TypeError("Cannot add/subtract non-tick DateOffset to {cls}" .format(cls=type(self).__name__)) - def _evaluate_with_timedelta_like(self, other, op): - if isinstance(other, ABCSeries): - # GH#19042 + def __mul__(self, other): + other = lib.item_from_zerodim(other) + + if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)): return NotImplemented - opstr = '__{opname}__'.format(opname=op.__name__).replace('__r', '__') - # allow division by a timedelta - if opstr in ['__div__', '__truediv__', '__floordiv__']: - if _is_convertible_to_td(other): - other = Timedelta(other) - if isna(other): - raise NotImplementedError( - "division by pd.NaT not implemented") - - i8 = self.asi8 - left, right = i8, other.value - - if opstr in ['__floordiv__']: - result = op(left, right) - else: - result = op(left, np.float64(right)) - result = self._maybe_mask_results(result, fill_value=None, - convert='float64') - return result + if is_scalar(other): + # numpy will accept float and int, raise TypeError for others + result = self._data * other + freq = None + if self.freq is not None and not isna(other): + freq = self.freq * other + return type(self)(result, freq=freq) + + if not hasattr(other, "dtype"): + # list, tuple + other = np.array(other) + if len(other) != len(self) and not is_timedelta64_dtype(other): + # Exclude timedelta64 here so we correctly raise TypeError + # for that instead of ValueError + raise ValueError("Cannot multiply with unequal lengths") + + if is_object_dtype(other): + # this multiplication will succeed only if all elements of other + # are int or float scalars, so we will end up with + # timedelta64[ns]-dtyped result + result = [self[n] * other[n] for n in range(len(self))] + result = np.array(result) + return type(self)(result) - return NotImplemented + # numpy will accept float or int dtype, raise TypeError for others + result = self._data * other + return type(self)(result) - __mul__ = _wrap_tdi_op(operator.mul) __rmul__ = __mul__ - __floordiv__ = _wrap_tdi_op(operator.floordiv) - __rfloordiv__ = _wrap_tdi_op(ops.rfloordiv) def __truediv__(self, other): # timedelta / X is well-defined for timedelta-like or numeric X @@ -464,6 +443,165 @@ def __rtruediv__(self, other): __div__ = __truediv__ __rdiv__ = __rtruediv__ + def __floordiv__(self, other): + if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)): + return NotImplemented + + other = lib.item_from_zerodim(other) + if is_scalar(other): + if isinstance(other, (timedelta, np.timedelta64, Tick)): + other = Timedelta(other) + if other is NaT: + # treat this specifically as timedelta-NaT + result = np.empty(self.shape, dtype=np.float64) + result.fill(np.nan) + return result + + # dispatch to Timedelta implementation + result = other.__rfloordiv__(self._data) + return result + + # at this point we should only have numeric scalars; anything + # else will raise + result = self.asi8 // other + result[self._isnan] = iNaT + freq = None + if self.freq is not None: + # Note: freq gets division, not floor-division + freq = self.freq / other + return type(self)(result.view('m8[ns]'), freq=freq) + + if not hasattr(other, "dtype"): + # list, tuple + other = np.array(other) + if len(other) != len(self): + raise ValueError("Cannot divide with unequal lengths") + + elif is_timedelta64_dtype(other): + other = type(self)(other) + + # numpy timedelta64 does not natively support floordiv, so operate + # on the i8 values + result = self.asi8 // other.asi8 + mask = self._isnan | other._isnan + if mask.any(): + result = result.astype(np.int64) + result[mask] = np.nan + return result + + elif is_object_dtype(other): + result = [self[n] // other[n] for n in range(len(self))] + result = np.array(result) + if lib.infer_dtype(result) == 'timedelta': + result, _ = sequence_to_td64ns(result) + return type(self)(result) + return result + + elif is_integer_dtype(other) or is_float_dtype(other): + result = self._data // other + return type(self)(result) + + else: + dtype = getattr(other, "dtype", type(other).__name__) + raise TypeError("Cannot divide {typ} by {cls}" + .format(typ=dtype, cls=type(self).__name__)) + + def __rfloordiv__(self, other): + if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)): + return NotImplemented + + other = lib.item_from_zerodim(other) + if is_scalar(other): + if isinstance(other, (timedelta, np.timedelta64, Tick)): + other = Timedelta(other) + if other is NaT: + # treat this specifically as timedelta-NaT + result = np.empty(self.shape, dtype=np.float64) + result.fill(np.nan) + return result + + # dispatch to Timedelta implementation + result = other.__floordiv__(self._data) + return result + + raise TypeError("Cannot divide {typ} by {cls}" + .format(typ=type(other).__name__, + cls=type(self).__name__)) + + if not hasattr(other, "dtype"): + # list, tuple + other = np.array(other) + if len(other) != len(self): + raise ValueError("Cannot divide with unequal lengths") + + elif is_timedelta64_dtype(other): + other = type(self)(other) + + # numpy timedelta64 does not natively support floordiv, so operate + # on the i8 values + result = other.asi8 // self.asi8 + mask = self._isnan | other._isnan + if mask.any(): + result = result.astype(np.int64) + result[mask] = np.nan + return result + + elif is_object_dtype(other): + result = [other[n] // self[n] for n in range(len(self))] + result = np.array(result) + return result + + else: + dtype = getattr(other, "dtype", type(other).__name__) + raise TypeError("Cannot divide {typ} by {cls}" + .format(typ=dtype, cls=type(self).__name__)) + + def __mod__(self, other): + # Note: This is a naive implementation, can likely be optimized + if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)): + return NotImplemented + + other = lib.item_from_zerodim(other) + if isinstance(other, (timedelta, np.timedelta64, Tick)): + other = Timedelta(other) + return self - (self // other) * other + + def __rmod__(self, other): + # Note: This is a naive implementation, can likely be optimized + if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)): + return NotImplemented + + other = lib.item_from_zerodim(other) + if isinstance(other, (timedelta, np.timedelta64, Tick)): + other = Timedelta(other) + return other - (other // self) * self + + def __divmod__(self, other): + # Note: This is a naive implementation, can likely be optimized + if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)): + return NotImplemented + + other = lib.item_from_zerodim(other) + if isinstance(other, (timedelta, np.timedelta64, Tick)): + other = Timedelta(other) + + res1 = self // other + res2 = self - res1 * other + return res1, res2 + + def __rdivmod__(self, other): + # Note: This is a naive implementation, can likely be optimized + if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)): + return NotImplemented + + other = lib.item_from_zerodim(other) + if isinstance(other, (timedelta, np.timedelta64, Tick)): + other = Timedelta(other) + + res1 = other // self + res2 = other - res1 * self + return res1, res2 + # Note: TimedeltaIndex overrides this in call to cls._add_numeric_methods def __neg__(self): if self.freq is not None: diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index e850db4178f41..a5b8e22070923 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -5031,23 +5031,22 @@ def _add_numeric_methods_binary(cls): cls.__radd__ = _make_arithmetic_op(ops.radd, cls) cls.__sub__ = _make_arithmetic_op(operator.sub, cls) cls.__rsub__ = _make_arithmetic_op(ops.rsub, cls) - cls.__mul__ = _make_arithmetic_op(operator.mul, cls) - cls.__rmul__ = _make_arithmetic_op(ops.rmul, cls) cls.__rpow__ = _make_arithmetic_op(ops.rpow, cls) cls.__pow__ = _make_arithmetic_op(operator.pow, cls) + + cls.__truediv__ = _make_arithmetic_op(operator.truediv, cls) + cls.__rtruediv__ = _make_arithmetic_op(ops.rtruediv, cls) + if not compat.PY3: + cls.__div__ = _make_arithmetic_op(operator.div, cls) + cls.__rdiv__ = _make_arithmetic_op(ops.rdiv, cls) + + # TODO: rmod? rdivmod? cls.__mod__ = _make_arithmetic_op(operator.mod, cls) cls.__floordiv__ = _make_arithmetic_op(operator.floordiv, cls) cls.__rfloordiv__ = _make_arithmetic_op(ops.rfloordiv, cls) - - if not issubclass(cls, ABCTimedeltaIndex): - # GH#23829 TimedeltaIndex defines these directly - cls.__truediv__ = _make_arithmetic_op(operator.truediv, cls) - cls.__rtruediv__ = _make_arithmetic_op(ops.rtruediv, cls) - if not compat.PY3: - cls.__div__ = _make_arithmetic_op(operator.div, cls) - cls.__rdiv__ = _make_arithmetic_op(ops.rdiv, cls) - cls.__divmod__ = _make_arithmetic_op(divmod, cls) + cls.__mul__ = _make_arithmetic_op(operator.mul, cls) + cls.__rmul__ = _make_arithmetic_op(ops.rmul, cls) @classmethod def _add_numeric_methods_unary(cls): diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index 0dedd8fe1cf4b..bb0d54f195457 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -573,6 +573,12 @@ def wrap_arithmetic_op(self, other, result): if result is NotImplemented: return NotImplemented + if isinstance(result, tuple): + # divmod, rdivmod + assert len(result) == 2 + return (wrap_arithmetic_op(self, other, result[0]), + wrap_arithmetic_op(self, other, result[1])) + if not isinstance(result, Index): # Index.__new__ will choose appropriate subclass for dtype result = Index(result) diff --git a/pandas/core/indexes/timedeltas.py b/pandas/core/indexes/timedeltas.py index e33d61d29d302..1c84e592d3a0d 100644 --- a/pandas/core/indexes/timedeltas.py +++ b/pandas/core/indexes/timedeltas.py @@ -31,7 +31,24 @@ from pandas.tseries.frequencies import to_offset -class TimedeltaIndex(TimedeltaArray, DatetimeIndexOpsMixin, Int64Index): +def _make_wrapped_arith_op(opname): + + meth = getattr(TimedeltaArray, opname) + + def method(self, other): + oth = other + if isinstance(other, Index): + oth = other._data + + result = meth(self, oth) + return wrap_arithmetic_op(self, other, result) + + method.__name__ = opname + return method + + +class TimedeltaIndex(TimedeltaArray, DatetimeIndexOpsMixin, + dtl.TimelikeOps, Int64Index): """ Immutable ndarray of timedelta64 data, represented internally as int64, and which can be boxed to timedelta objects @@ -203,10 +220,6 @@ def _maybe_update_attributes(self, attrs): attrs['freq'] = 'infer' return attrs - def _evaluate_with_timedelta_like(self, other, op): - result = TimedeltaArray._evaluate_with_timedelta_like(self, other, op) - return wrap_arithmetic_op(self, other, result) - # ------------------------------------------------------------------- # Rendering Methods @@ -224,10 +237,14 @@ def _format_native_types(self, na_rep=u'NaT', date_format=None, **kwargs): # ------------------------------------------------------------------- # Wrapping TimedeltaArray - __mul__ = Index.__mul__ - __rmul__ = Index.__rmul__ - __floordiv__ = Index.__floordiv__ - __rfloordiv__ = Index.__rfloordiv__ + __mul__ = _make_wrapped_arith_op("__mul__") + __rmul__ = _make_wrapped_arith_op("__rmul__") + __floordiv__ = _make_wrapped_arith_op("__floordiv__") + __rfloordiv__ = _make_wrapped_arith_op("__rfloordiv__") + __mod__ = _make_wrapped_arith_op("__mod__") + __rmod__ = _make_wrapped_arith_op("__rmod__") + __divmod__ = _make_wrapped_arith_op("__divmod__") + __rdivmod__ = _make_wrapped_arith_op("__rdivmod__") days = wrap_field_accessor(TimedeltaArray.days) seconds = wrap_field_accessor(TimedeltaArray.seconds) @@ -658,7 +675,7 @@ def delete(self, loc): TimedeltaIndex._add_comparison_ops() -TimedeltaIndex._add_numeric_methods() +TimedeltaIndex._add_numeric_methods_unary() TimedeltaIndex._add_logical_methods_disabled() TimedeltaIndex._add_datetimelike_methods() diff --git a/pandas/core/ops.py b/pandas/core/ops.py index 850c4dd7c45e7..6ea31422478f2 100644 --- a/pandas/core/ops.py +++ b/pandas/core/ops.py @@ -1552,8 +1552,7 @@ def wrapper(left, right): elif is_timedelta64_dtype(left): result = dispatch_to_index_op(op, left, right, pd.TimedeltaIndex) return construct_result(left, result, - index=left.index, name=res_name, - dtype=result.dtype) + index=left.index, name=res_name) elif is_timedelta64_dtype(right): # We should only get here with non-scalar or timedelta64('NaT') diff --git a/pandas/tests/arithmetic/test_datetime64.py b/pandas/tests/arithmetic/test_datetime64.py index 667c2b4103e00..bc9b712e78d03 100644 --- a/pandas/tests/arithmetic/test_datetime64.py +++ b/pandas/tests/arithmetic/test_datetime64.py @@ -1413,7 +1413,7 @@ def check(get_ser, test_ser): # with 'operate' (from core/ops.py) for the ops that are not # defined op = getattr(get_ser, op_str, None) - with pytest.raises(TypeError, match='operate|cannot'): + with pytest.raises(TypeError, match='operate|[cC]annot'): op(test_ser) # ## timedelta64 ### diff --git a/pandas/tests/arithmetic/test_timedelta64.py b/pandas/tests/arithmetic/test_timedelta64.py index 81e7062c23fbe..5f2fd98e29b96 100644 --- a/pandas/tests/arithmetic/test_timedelta64.py +++ b/pandas/tests/arithmetic/test_timedelta64.py @@ -1280,7 +1280,8 @@ def test_td64arr_floordiv_int(self, box_with_array): result = idx // 1 tm.assert_equal(result, idx) - pattern = 'floor_divide cannot use operands' + pattern = ('floor_divide cannot use operands|' + 'Cannot divide int by Timedelta*') with pytest.raises(TypeError, match=pattern): 1 // idx @@ -1317,6 +1318,66 @@ def test_td64arr_rfloordiv_tdlike_scalar(self, scalar_td, box_with_array): res = tdi // (scalar_td) tm.assert_equal(res, expected) + # ------------------------------------------------------------------ + # mod, divmod + # TODO: operations with timedelta-like arrays, numeric arrays, + # reversed ops + + def test_td64arr_mod_tdscalar(self, box_with_array, three_days): + tdi = timedelta_range('1 Day', '9 days') + tdarr = tm.box_expected(tdi, box_with_array) + + expected = TimedeltaIndex(['1 Day', '2 Days', '0 Days'] * 3) + expected = tm.box_expected(expected, box_with_array) + + result = tdarr % three_days + tm.assert_equal(result, expected) + + if box_with_array is pd.DataFrame: + pytest.xfail("DataFrame does not have __divmod__ or __rdivmod__") + + result = divmod(tdarr, three_days) + tm.assert_equal(result[1], expected) + tm.assert_equal(result[0], tdarr // three_days) + + def test_td64arr_mod_int(self, box_with_array): + tdi = timedelta_range('1 ns', '10 ns', periods=10) + tdarr = tm.box_expected(tdi, box_with_array) + + expected = TimedeltaIndex(['1 ns', '0 ns'] * 5) + expected = tm.box_expected(expected, box_with_array) + + result = tdarr % 2 + tm.assert_equal(result, expected) + + with pytest.raises(TypeError): + 2 % tdarr + + if box_with_array is pd.DataFrame: + pytest.xfail("DataFrame does not have __divmod__ or __rdivmod__") + + result = divmod(tdarr, 2) + tm.assert_equal(result[1], expected) + tm.assert_equal(result[0], tdarr // 2) + + def test_td64arr_rmod_tdscalar(self, box_with_array, three_days): + tdi = timedelta_range('1 Day', '9 days') + tdarr = tm.box_expected(tdi, box_with_array) + + expected = ['0 Days', '1 Day', '0 Days'] + ['3 Days'] * 6 + expected = TimedeltaIndex(expected) + expected = tm.box_expected(expected, box_with_array) + + result = three_days % tdarr + tm.assert_equal(result, expected) + + if box_with_array is pd.DataFrame: + pytest.xfail("DataFrame does not have __divmod__ or __rdivmod__") + + result = divmod(three_days, tdarr) + tm.assert_equal(result[1], expected) + tm.assert_equal(result[0], three_days // tdarr) + # ------------------------------------------------------------------ # Operations with invalid others