From 2cbde5184ed7948e132f0fb7a5a533410bdc7be4 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Wed, 22 Jan 2020 18:18:30 -0800 Subject: [PATCH] WIP/TST: add dt64tz to indices fixture --- pandas/core/indexes/datetimelike.py | 76 ++++++++++++++++++++++- pandas/core/series.py | 2 +- pandas/tests/arithmetic/test_period.py | 1 + pandas/tests/indexes/common.py | 10 +++ pandas/tests/indexes/conftest.py | 1 + pandas/tests/indexes/test_numpy_compat.py | 8 ++- 6 files changed, 94 insertions(+), 4 deletions(-) diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index 1bfec9fbad0ed..753e47387d448 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -23,13 +23,14 @@ is_list_like, is_period_dtype, is_scalar, + is_timedelta64_dtype, needs_i8_conversion, ) from pandas.core.dtypes.concat import concat_compat from pandas.core.dtypes.generic import ABCIndex, ABCIndexClass, ABCSeries from pandas.core.dtypes.missing import is_valid_nat_for_dtype, isna -from pandas.core import algorithms +from pandas.core import algorithms, ops from pandas.core.accessor import PandasDelegate from pandas.core.arrays import DatetimeArray, ExtensionArray, TimedeltaArray from pandas.core.arrays.datetimelike import DatetimeLikeArrayMixin @@ -126,6 +127,79 @@ def __array_wrap__(self, result, context=None): attrs["freq"] = "infer" return Index(result, **attrs) + def __array_ufunc__(self, ufunc, method: str, *inputs, **kwargs): + result = ops.maybe_dispatch_ufunc_to_dunder_op( + self, ufunc, method, *inputs, **kwargs + ) + if result is not NotImplemented: + return result + + # unary ufuncs + if ( + len(inputs) == 1 + and inputs[0] is self + and not kwargs + and method == "__call__" + ): + if ufunc.__name__ in ["isfinite"]: + return self.notna() + if ufunc.__name__ == "negative": + return -self + if ufunc.__name__ == "absolute": + if is_timedelta64_dtype(self.dtype): + return abs(self) + + # binary ufuncs + if len(inputs) == 2 and not kwargs and method == "__call__": + from pandas import DatetimeIndex, TimedeltaIndex + + new_inputs = [inputs[0], inputs[1]] + for i, val in enumerate(new_inputs): + if isinstance(val, np.ndarray) and val.dtype.kind == "m": + val = TimedeltaIndex(val) + new_inputs[i] = val + if isinstance(val, np.ndarray) and val.dtype.kind == "M": + val = DatetimeIndex(val) + new_inputs[i] = val + + if ufunc.__name__ == "subtract": + if not isinstance(inputs[0], DatetimeIndexOpsMixin): + result = new_inputs[1].__rsub__(new_inputs[0]) + else: + result = new_inputs[0].__sub__(new_inputs[1]) + + if result is NotImplemented: + # raise explicitly or else we get recursionerror + raise TypeError("cannot use operands with types") + return result + + if ufunc.__name__ == "add": + if not isinstance(inputs[0], DatetimeIndexOpsMixin): + result = new_inputs[1].__radd__(new_inputs[0]) + else: + result = new_inputs[0].__add__(new_inputs[1]) + + if result is NotImplemented: + # raise explicitly or else we get recursionerror + raise TypeError("cannot use operands with types") + return result + + if ( + len(inputs) == 2 + and len(kwargs) == 1 + and "out" in kwargs + and method == "__call__" + ): + out = kwargs["out"] + if ufunc.__name__ == "subtract": + out[:] = inputs[0] - inputs[1] + return + if ufunc.__name__ == "add": + out[:] = inputs[0] + inputs[1] + return + + return NotImplemented + # ------------------------------------------------------------------------ def equals(self, other) -> bool: diff --git a/pandas/core/series.py b/pandas/core/series.py index ffe0642f799fa..f46ecf28d65c2 100644 --- a/pandas/core/series.py +++ b/pandas/core/series.py @@ -660,7 +660,7 @@ def __array_ufunc__( result = getattr(ufunc, method)(*inputs, **kwargs) name: Optional[Hashable] - if len(set(names)) == 1: + if len(set(names)) == 1: # TODO: we have a func in ops for this name = names[0] else: name = None diff --git a/pandas/tests/arithmetic/test_period.py b/pandas/tests/arithmetic/test_period.py index abb667260f094..afc974dc18912 100644 --- a/pandas/tests/arithmetic/test_period.py +++ b/pandas/tests/arithmetic/test_period.py @@ -1278,6 +1278,7 @@ def test_parr_ops_errors(self, ng, func, box_with_array): ) obj = tm.box_expected(idx, box_with_array) msg = ( + "cannot use operands with types|" r"unsupported operand type\(s\)|can only concatenate|" r"must be str|object to str implicitly" ) diff --git a/pandas/tests/indexes/common.py b/pandas/tests/indexes/common.py index f3ebe8313d0c6..17d2b85861468 100644 --- a/pandas/tests/indexes/common.py +++ b/pandas/tests/indexes/common.py @@ -6,6 +6,7 @@ from pandas._libs.tslib import iNaT +from pandas.core.dtypes.common import is_datetime64tz_dtype from pandas.core.dtypes.dtypes import CategoricalDtype import pandas as pd @@ -301,6 +302,9 @@ def test_ensure_copied_data(self, indices): index_type = type(indices) result = index_type(indices.values, copy=True, **init_kwargs) + if isinstance(indices, DatetimeIndex) and indices.tz is not None: + result = result.tz_localize("UTC").tz_convert(indices.tz) + tm.assert_index_equal(indices, result) tm.assert_numpy_array_equal( indices._ndarray_values, result._ndarray_values, check_same="copy" @@ -467,6 +471,9 @@ def test_intersection_base(self, indices): # GH 10149 cases = [klass(second.values) for klass in [np.array, Series, list]] for case in cases: + if is_datetime64tz_dtype(indices.dtype): + # the second.values will drop tz, so this would be wrong + continue result = first.intersection(case) assert tm.equalContents(result, second) @@ -485,6 +492,9 @@ def test_union_base(self, indices): # GH 10149 cases = [klass(second.values) for klass in [np.array, Series, list]] for case in cases: + if is_datetime64tz_dtype(indices.dtype): + # the second.values will drop tz, so this would be wrong + continue if not isinstance(indices, CategoricalIndex): result = first.union(case) assert tm.equalContents(result, everything) diff --git a/pandas/tests/indexes/conftest.py b/pandas/tests/indexes/conftest.py index e3e7ff4093b76..6828388dce901 100644 --- a/pandas/tests/indexes/conftest.py +++ b/pandas/tests/indexes/conftest.py @@ -9,6 +9,7 @@ "unicode": tm.makeUnicodeIndex(100), "string": tm.makeStringIndex(100), "datetime": tm.makeDateIndex(100), + "datetime-tz": tm.makeDateIndex(100, tz="US/Pacific"), "period": tm.makePeriodIndex(100), "timedelta": tm.makeTimedeltaIndex(100), "int": tm.makeIntIndex(100), diff --git a/pandas/tests/indexes/test_numpy_compat.py b/pandas/tests/indexes/test_numpy_compat.py index 583556656ac87..2a8798b304ed7 100644 --- a/pandas/tests/indexes/test_numpy_compat.py +++ b/pandas/tests/indexes/test_numpy_compat.py @@ -99,8 +99,12 @@ def test_numpy_ufuncs_other(indices, func): elif isinstance(idx, PeriodIndex): # raise TypeError or ValueError (PeriodIndex) - with pytest.raises(Exception): - func(idx) + if func in [np.isfinite]: + result = func(idx) + assert isinstance(result, np.ndarray) + else: + with pytest.raises(Exception): + func(idx) elif isinstance(idx, (Float64Index, Int64Index, UInt64Index)): # Results in bool array