diff --git a/pandas/core/arrays/numpy_.py b/pandas/core/arrays/numpy_.py index 237d571507a3a..05139783456b9 100644 --- a/pandas/core/arrays/numpy_.py +++ b/pandas/core/arrays/numpy_.py @@ -362,19 +362,29 @@ def __invert__(self): @classmethod def _create_arithmetic_method(cls, op): + + pd_op = ops.get_array_op(op) + @ops.unpack_zerodim_and_defer(op.__name__) def arithmetic_method(self, other): if isinstance(other, cls): other = other._ndarray - with np.errstate(all="ignore"): - result = op(self._ndarray, other) + result = pd_op(self._ndarray, other) - if op is divmod: + if op is divmod or op is ops.rdivmod: a, b = result - return cls(a), cls(b) - - return cls(result) + if isinstance(a, np.ndarray): + # for e.g. op vs TimedeltaArray, we may already + # have an ExtensionArray, in which case we do not wrap + return cls(a), cls(b) + return a, b + + if isinstance(result, np.ndarray): + # for e.g. multiplication vs TimedeltaArray, we may already + # have an ExtensionArray, in which case we do not wrap + return cls(result) + return result return compat.set_function_name(arithmetic_method, f"__{op.__name__}__", cls) diff --git a/pandas/tests/arithmetic/common.py b/pandas/tests/arithmetic/common.py index cd8dd102dc27c..a663c2f3a0175 100644 --- a/pandas/tests/arithmetic/common.py +++ b/pandas/tests/arithmetic/common.py @@ -4,7 +4,7 @@ import numpy as np import pytest -from pandas import DataFrame, Index, Series +from pandas import DataFrame, Index, Series, array as pd_array import pandas._testing as tm @@ -49,12 +49,12 @@ def assert_invalid_comparison(left, right, box): ---------- left : np.ndarray, ExtensionArray, Index, or Series right : object - box : {pd.DataFrame, pd.Series, pd.Index, tm.to_array} + box : {pd.DataFrame, pd.Series, pd.Index, pd.array, tm.to_array} """ # Not for tznaive-tzaware comparison # Note: not quite the same as how we do this for tm.box_expected - xbox = box if box is not Index else np.array + xbox = box if box not in [Index, pd_array] else np.array result = left == right expected = xbox(np.zeros(result.shape, dtype=np.bool_)) diff --git a/pandas/tests/arithmetic/conftest.py b/pandas/tests/arithmetic/conftest.py index 8b9e5cd371a90..6286711ac6113 100644 --- a/pandas/tests/arithmetic/conftest.py +++ b/pandas/tests/arithmetic/conftest.py @@ -2,7 +2,6 @@ import pytest import pandas as pd -import pandas._testing as tm # ------------------------------------------------------------------ # Helper Functions @@ -56,7 +55,7 @@ def one(request): zeros = [ box_cls([0] * 5, dtype=dtype) - for box_cls in [pd.Index, np.array] + for box_cls in [pd.Index, np.array, pd.array] for dtype in [np.int64, np.uint64, np.float64] ] zeros.extend( @@ -231,7 +230,7 @@ def box(request): return request.param -@pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, tm.to_array], ids=id_func) +@pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, pd.array], ids=id_func) def box_with_array(request): """ Fixture to test behavior for Index, Series, DataFrame, and pandas Array diff --git a/pandas/tests/arithmetic/test_datetime64.py b/pandas/tests/arithmetic/test_datetime64.py index 0dd389ed516c7..626dd4f748e0b 100644 --- a/pandas/tests/arithmetic/test_datetime64.py +++ b/pandas/tests/arithmetic/test_datetime64.py @@ -48,7 +48,9 @@ def test_compare_zerodim(self, tz_naive_fixture, box_with_array): # Test comparison with zero-dimensional array is unboxed tz = tz_naive_fixture box = box_with_array - xbox = box_with_array if box_with_array is not pd.Index else np.ndarray + xbox = ( + box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray + ) dti = date_range("20130101", periods=3, tz=tz) other = np.array(dti.to_numpy()[0]) @@ -135,7 +137,7 @@ def test_dt64arr_nat_comparison(self, tz_naive_fixture, box_with_array): # GH#22242, GH#22163 DataFrame considered NaT == ts incorrectly tz = tz_naive_fixture box = box_with_array - xbox = box if box is not pd.Index else np.ndarray + xbox = box if box not in [pd.Index, pd.array] else np.ndarray ts = pd.Timestamp.now(tz) ser = pd.Series([ts, pd.NaT]) @@ -203,6 +205,8 @@ def test_nat_comparisons(self, dtype, index_or_series, reverse, pair): def test_comparison_invalid(self, tz_naive_fixture, box_with_array): # GH#4968 # invalid date/int comparisons + if box_with_array is pd.array: + pytest.xfail("assert_invalid_comparison doesnt handle BooleanArray yet") tz = tz_naive_fixture ser = Series(range(5)) ser2 = Series(pd.date_range("20010101", periods=5, tz=tz)) @@ -226,8 +230,12 @@ def test_nat_comparisons_scalar(self, dtype, data, box_with_array): # dont bother testing ndarray comparison methods as this fails # on older numpys (since they check object identity) return + if box_with_array is pd.array and dtype is object: + pytest.xfail("reversed comparisons give BooleanArray, not ndarray") - xbox = box_with_array if box_with_array is not pd.Index else np.ndarray + xbox = ( + box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray + ) left = Series(data, dtype=dtype) left = tm.box_expected(left, box_with_array) @@ -299,7 +307,9 @@ def test_timestamp_compare_series(self, left, right): def test_dt64arr_timestamp_equality(self, box_with_array): # GH#11034 - xbox = box_with_array if box_with_array is not pd.Index else np.ndarray + xbox = ( + box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray + ) ser = pd.Series([pd.Timestamp("2000-01-29 01:59:00"), "NaT"]) ser = tm.box_expected(ser, box_with_array) @@ -388,7 +398,9 @@ def test_dti_cmp_nat(self, dtype, box_with_array): # on older numpys (since they check object identity) return - xbox = box_with_array if box_with_array is not pd.Index else np.ndarray + xbox = ( + box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray + ) left = pd.DatetimeIndex( [pd.Timestamp("2011-01-01"), pd.NaT, pd.Timestamp("2011-01-03")] diff --git a/pandas/tests/arithmetic/test_numeric.py b/pandas/tests/arithmetic/test_numeric.py index 139401bdf5806..df98b43e11f4a 100644 --- a/pandas/tests/arithmetic/test_numeric.py +++ b/pandas/tests/arithmetic/test_numeric.py @@ -89,8 +89,9 @@ def test_compare_invalid(self): b.name = pd.Timestamp("2000-01-01") tm.assert_series_equal(a / b, 1 / (b / a)) - def test_numeric_cmp_string_numexpr_path(self, box): + def test_numeric_cmp_string_numexpr_path(self, box_with_array): # GH#36377, GH#35700 + box = box_with_array xbox = box if box is not pd.Index else np.ndarray obj = pd.Series(np.random.randn(10 ** 5)) @@ -183,10 +184,14 @@ def test_ops_series(self): ], ids=lambda x: type(x).__name__, ) - def test_numeric_arr_mul_tdscalar(self, scalar_td, numeric_idx, box): + def test_numeric_arr_mul_tdscalar(self, scalar_td, numeric_idx, box_with_array): # GH#19333 + box = box_with_array + if box is pd.array: + pytest.xfail( + "we get a PandasArray[timedelta64[ns]] instead of TimedeltaArray" + ) index = numeric_idx - expected = pd.TimedeltaIndex([pd.Timedelta(days=n) for n in range(5)]) index = tm.box_expected(index, box) @@ -207,7 +212,11 @@ def test_numeric_arr_mul_tdscalar(self, scalar_td, numeric_idx, box): ], ids=lambda x: type(x).__name__, ) - def test_numeric_arr_mul_tdscalar_numexpr_path(self, scalar_td, box): + def test_numeric_arr_mul_tdscalar_numexpr_path(self, scalar_td, box_with_array): + box = box_with_array + if box is pd.array: + pytest.xfail("IntegerArray.__mul__ doesnt handle timedeltas") + arr = np.arange(2 * 10 ** 4).astype(np.int64) obj = tm.box_expected(arr, box, transpose=False) @@ -220,7 +229,11 @@ def test_numeric_arr_mul_tdscalar_numexpr_path(self, scalar_td, box): result = scalar_td * obj tm.assert_equal(result, expected) - def test_numeric_arr_rdiv_tdscalar(self, three_days, numeric_idx, box): + def test_numeric_arr_rdiv_tdscalar(self, three_days, numeric_idx, box_with_array): + box = box_with_array + if box is pd.array: + pytest.xfail("We get PandasArray[td64] instead of TimedeltaArray") + index = numeric_idx[1:3] expected = TimedeltaIndex(["3 Days", "36 Hours"]) @@ -248,7 +261,11 @@ def test_numeric_arr_rdiv_tdscalar(self, three_days, numeric_idx, box): pd.offsets.Second(0), ], ) - def test_add_sub_timedeltalike_invalid(self, numeric_idx, other, box): + def test_add_sub_timedeltalike_invalid(self, numeric_idx, other, box_with_array): + box = box_with_array + if box is pd.array: + pytest.xfail("PandasArray[int].__add__ doesnt raise on td64") + left = tm.box_expected(numeric_idx, box) msg = ( "unsupported operand type|" @@ -276,16 +293,21 @@ def test_add_sub_timedeltalike_invalid(self, numeric_idx, other, box): ], ) @pytest.mark.filterwarnings("ignore:elementwise comp:DeprecationWarning") - def test_add_sub_datetimelike_invalid(self, numeric_idx, other, box): + def test_add_sub_datetimelike_invalid(self, numeric_idx, other, box_with_array): # GH#28080 numeric+datetime64 should raise; Timestamp raises # NullFrequencyError instead of TypeError so is excluded. + box = box_with_array left = tm.box_expected(numeric_idx, box) - msg = ( - "unsupported operand type|" - "Cannot (add|subtract) NaT (to|from) ndarray|" - "Addition/subtraction of integers and integer-arrays|" - "Concatenation operation is not implemented for NumPy arrays" + msg = "|".join( + [ + "unsupported operand type", + "Cannot (add|subtract) NaT (to|from) ndarray", + "Addition/subtraction of integers and integer-arrays", + "Concatenation operation is not implemented for NumPy arrays", + # pd.array vs np.datetime64 case + r"operand type\(s\) all returned NotImplemented from __array_ufunc__", + ] ) with pytest.raises(TypeError, match=msg): left + other @@ -568,8 +590,9 @@ class TestMultiplicationDivision: # __mul__, __rmul__, __div__, __rdiv__, __floordiv__, __rfloordiv__ # for non-timestamp/timedelta/period dtypes - def test_divide_decimal(self, box): + def test_divide_decimal(self, box_with_array): # resolves issue GH#9787 + box = box_with_array ser = Series([Decimal(10)]) expected = Series([Decimal(5)]) diff --git a/pandas/tests/arithmetic/test_object.py b/pandas/tests/arithmetic/test_object.py index c0cb522b516ab..02cb4f4d7a606 100644 --- a/pandas/tests/arithmetic/test_object.py +++ b/pandas/tests/arithmetic/test_object.py @@ -104,22 +104,22 @@ def test_add_extension_scalar(self, other, box_with_array, op): result = op(arr, other) tm.assert_equal(result, expected) - def test_objarr_add_str(self, box): + def test_objarr_add_str(self, box_with_array): ser = pd.Series(["x", np.nan, "x"]) expected = pd.Series(["xa", np.nan, "xa"]) - ser = tm.box_expected(ser, box) - expected = tm.box_expected(expected, box) + ser = tm.box_expected(ser, box_with_array) + expected = tm.box_expected(expected, box_with_array) result = ser + "a" tm.assert_equal(result, expected) - def test_objarr_radd_str(self, box): + def test_objarr_radd_str(self, box_with_array): ser = pd.Series(["x", np.nan, "x"]) expected = pd.Series(["ax", np.nan, "ax"]) - ser = tm.box_expected(ser, box) - expected = tm.box_expected(expected, box) + ser = tm.box_expected(ser, box_with_array) + expected = tm.box_expected(expected, box_with_array) result = "a" + ser tm.assert_equal(result, expected) diff --git a/pandas/tests/arithmetic/test_period.py b/pandas/tests/arithmetic/test_period.py index 930435074efc1..e78e696d00398 100644 --- a/pandas/tests/arithmetic/test_period.py +++ b/pandas/tests/arithmetic/test_period.py @@ -28,7 +28,9 @@ class TestPeriodArrayLikeComparisons: def test_compare_zerodim(self, box_with_array): # GH#26689 make sure we unbox zero-dimensional arrays - xbox = box_with_array if box_with_array is not pd.Index else np.ndarray + xbox = ( + box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray + ) pi = pd.period_range("2000", periods=4) other = np.array(pi.to_numpy()[0]) @@ -68,7 +70,7 @@ def test_compare_object_dtype(self, box_with_array, other_box): pi = pd.period_range("2000", periods=5) parr = tm.box_expected(pi, box_with_array) - xbox = np.ndarray if box_with_array is pd.Index else box_with_array + xbox = np.ndarray if box_with_array in [pd.Index, pd.array] else box_with_array other = other_box(pi) @@ -175,7 +177,9 @@ def test_pi_cmp_period(self): # TODO: moved from test_datetime64; de-duplicate with version below def test_parr_cmp_period_scalar2(self, box_with_array): - xbox = box_with_array if box_with_array is not pd.Index else np.ndarray + xbox = ( + box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray + ) pi = pd.period_range("2000-01-01", periods=10, freq="D") @@ -196,7 +200,7 @@ def test_parr_cmp_period_scalar2(self, box_with_array): @pytest.mark.parametrize("freq", ["M", "2M", "3M"]) def test_parr_cmp_period_scalar(self, freq, box_with_array): # GH#13200 - xbox = np.ndarray if box_with_array is pd.Index else box_with_array + xbox = np.ndarray if box_with_array in [pd.Index, pd.array] else box_with_array base = PeriodIndex(["2011-01", "2011-02", "2011-03", "2011-04"], freq=freq) base = tm.box_expected(base, box_with_array) @@ -235,7 +239,7 @@ def test_parr_cmp_period_scalar(self, freq, box_with_array): @pytest.mark.parametrize("freq", ["M", "2M", "3M"]) def test_parr_cmp_pi(self, freq, box_with_array): # GH#13200 - xbox = np.ndarray if box_with_array is pd.Index else box_with_array + xbox = np.ndarray if box_with_array in [pd.Index, pd.array] else box_with_array base = PeriodIndex(["2011-01", "2011-02", "2011-03", "2011-04"], freq=freq) base = tm.box_expected(base, box_with_array) @@ -284,7 +288,7 @@ def test_parr_cmp_pi_mismatched_freq_raises(self, freq, box_with_array): # TODO: Could parametrize over boxes for idx? idx = PeriodIndex(["2011", "2012", "2013", "2014"], freq="A") rev_msg = r"Input has different freq=(M|2M|3M) from PeriodArray\(freq=A-DEC\)" - idx_msg = rev_msg if box_with_array is tm.to_array else msg + idx_msg = rev_msg if box_with_array in [tm.to_array, pd.array] else msg with pytest.raises(IncompatibleFrequency, match=idx_msg): base <= idx @@ -298,7 +302,7 @@ def test_parr_cmp_pi_mismatched_freq_raises(self, freq, box_with_array): idx = PeriodIndex(["2011", "2012", "2013", "2014"], freq="4M") rev_msg = r"Input has different freq=(M|2M|3M) from PeriodArray\(freq=4M\)" - idx_msg = rev_msg if box_with_array is tm.to_array else msg + idx_msg = rev_msg if box_with_array in [tm.to_array, pd.array] else msg with pytest.raises(IncompatibleFrequency, match=idx_msg): base <= idx @@ -779,7 +783,7 @@ def test_pi_add_sub_td64_array_tick(self): @pytest.mark.parametrize("tdi_freq", [None, "H"]) def test_parr_sub_td64array(self, box_with_array, tdi_freq, pi_freq): box = box_with_array - xbox = box if box is not tm.to_array else pd.Index + xbox = box if box not in [pd.array, tm.to_array] else pd.Index tdi = TimedeltaIndex(["1 hours", "2 hours"], freq=tdi_freq) dti = Timestamp("2018-03-07 17:16:40") + tdi diff --git a/pandas/tests/arithmetic/test_timedelta64.py b/pandas/tests/arithmetic/test_timedelta64.py index 68bedcc099a91..b3dfb5d015ab4 100644 --- a/pandas/tests/arithmetic/test_timedelta64.py +++ b/pandas/tests/arithmetic/test_timedelta64.py @@ -50,7 +50,9 @@ class TestTimedelta64ArrayLikeComparisons: def test_compare_timedelta64_zerodim(self, box_with_array): # GH#26689 should unbox when comparing with zerodim array box = box_with_array - xbox = box_with_array if box_with_array is not pd.Index else np.ndarray + xbox = ( + box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray + ) tdi = pd.timedelta_range("2H", periods=4) other = np.array(tdi.to_numpy()[0]) @@ -73,7 +75,8 @@ def test_compare_timedelta64_zerodim(self, box_with_array): def test_compare_timedeltalike_scalar(self, box_with_array, td_scalar): # regression test for GH#5963 box = box_with_array - xbox = box if box is not pd.Index else np.ndarray + xbox = box if box not in [pd.Index, pd.array] else np.ndarray + ser = pd.Series([timedelta(days=1), timedelta(days=2)]) ser = tm.box_expected(ser, box) actual = ser > td_scalar @@ -85,6 +88,7 @@ def test_compare_timedeltalike_scalar(self, box_with_array, td_scalar): def test_td64_comparisons_invalid(self, box_with_array, invalid): # GH#13624 for str box = box_with_array + rng = timedelta_range("1 days", periods=10) obj = tm.box_expected(rng, box) @@ -1142,19 +1146,24 @@ def test_td64arr_add_sub_integer_array(self, box_with_array): # GH#19959, deprecated GH#22535 # GH#22696 for DataFrame case, check that we don't dispatch to numpy # implementation, which treats int64 as m8[ns] + box = box_with_array + xbox = np.ndarray if box is pd.array else box rng = timedelta_range("1 days 09:00:00", freq="H", periods=3) - tdarr = tm.box_expected(rng, box_with_array) - other = tm.box_expected([4, 3, 2], box_with_array) + tdarr = tm.box_expected(rng, box) + other = tm.box_expected([4, 3, 2], xbox) msg = "Addition/subtraction of integers and integer-arrays" assert_invalid_addsub_type(tdarr, other, msg) def test_td64arr_addsub_integer_array_no_freq(self, box_with_array): # GH#19959 + box = box_with_array + xbox = np.ndarray if box is pd.array else box + tdi = TimedeltaIndex(["1 Day", "NaT", "3 Hours"]) - tdarr = tm.box_expected(tdi, box_with_array) - other = tm.box_expected([14, -1, 16], box_with_array) + tdarr = tm.box_expected(tdi, box) + other = tm.box_expected([14, -1, 16], xbox) msg = "Addition/subtraction of integers" assert_invalid_addsub_type(tdarr, other, msg) @@ -1204,7 +1213,7 @@ def test_td64arr_add_sub_tdi(self, box_with_array, names): ) tdi = TimedeltaIndex(["0 days", "1 day"], name=names[0]) - tdi = np.array(tdi) if box is tm.to_array else tdi + tdi = np.array(tdi) if box in [tm.to_array, pd.array] else tdi ser = Series([Timedelta(hours=3), Timedelta(hours=4)], name=names[1]) expected = Series( [Timedelta(hours=3), Timedelta(days=1, hours=4)], name=names[2] @@ -1311,7 +1320,7 @@ def test_td64arr_add_offset_index(self, names, box_with_array): tdi = TimedeltaIndex(["1 days 00:00:00", "3 days 04:00:00"], name=names[0]) other = pd.Index([pd.offsets.Hour(n=1), pd.offsets.Minute(n=-2)], name=names[1]) - other = np.array(other) if box is tm.to_array else other + other = np.array(other) if box in [tm.to_array, pd.array] else other expected = TimedeltaIndex( [tdi[n] + other[n] for n in range(len(tdi))], freq="infer", name=names[2] @@ -1353,8 +1362,8 @@ def test_td64arr_add_offset_array(self, box_with_array): def test_td64arr_sub_offset_index(self, names, box_with_array): # GH#18824, GH#19744 box = box_with_array - xbox = box if box is not tm.to_array else pd.Index - exname = names[2] if box is not tm.to_array else names[1] + xbox = box if box not in [tm.to_array, pd.array] else pd.Index + exname = names[2] if box not in [tm.to_array, pd.array] else names[1] if box is pd.DataFrame and names[1] != names[0]: pytest.skip( @@ -1395,13 +1404,13 @@ def test_td64arr_sub_offset_array(self, box_with_array): def test_td64arr_with_offset_series(self, names, box_with_array): # GH#18849 box = box_with_array - box2 = Series if box in [pd.Index, tm.to_array] else box + box2 = Series if box in [pd.Index, tm.to_array, pd.array] else box if box is pd.DataFrame: # Since we are operating with a DataFrame and a non-DataFrame, # the non-DataFrame is cast to Series and its name ignored. exname = names[0] - elif box is tm.to_array: + elif box in [tm.to_array, pd.array]: exname = names[1] else: exname = names[2] @@ -1456,8 +1465,11 @@ def test_td64arr_addsub_anchored_offset_arraylike(self, obox, box_with_array): # Unsorted def test_td64arr_add_sub_object_array(self, box_with_array): + box = box_with_array + xbox = np.ndarray if box is pd.array else box + tdi = pd.timedelta_range("1 day", periods=3, freq="D") - tdarr = tm.box_expected(tdi, box_with_array) + tdarr = tm.box_expected(tdi, box) other = np.array( [pd.Timedelta(days=1), pd.offsets.Day(2), pd.Timestamp("2000-01-04")] @@ -1469,7 +1481,7 @@ def test_td64arr_add_sub_object_array(self, box_with_array): expected = pd.Index( [pd.Timedelta(days=2), pd.Timedelta(days=4), pd.Timestamp("2000-01-07")] ) - expected = tm.box_expected(expected, box_with_array) + expected = tm.box_expected(expected, xbox) tm.assert_equal(result, expected) msg = "unsupported operand type|cannot subtract a datelike" @@ -1483,7 +1495,7 @@ def test_td64arr_add_sub_object_array(self, box_with_array): expected = pd.Index( [pd.Timedelta(0), pd.Timedelta(0), pd.Timestamp("2000-01-01")] ) - expected = tm.box_expected(expected, box_with_array) + expected = tm.box_expected(expected, xbox) tm.assert_equal(result, expected) @@ -1536,7 +1548,7 @@ def test_tdi_mul_int_array(self, box_with_array): def test_tdi_mul_int_series(self, box_with_array): box = box_with_array - xbox = pd.Series if box in [pd.Index, tm.to_array] else box + xbox = pd.Series if box in [pd.Index, tm.to_array, pd.array] else box idx = TimedeltaIndex(np.arange(5, dtype="int64")) expected = TimedeltaIndex(np.arange(5, dtype="int64") ** 2) @@ -1549,7 +1561,7 @@ def test_tdi_mul_int_series(self, box_with_array): def test_tdi_mul_float_series(self, box_with_array): box = box_with_array - xbox = pd.Series if box in [pd.Index, tm.to_array] else box + xbox = pd.Series if box in [pd.Index, tm.to_array, pd.array] else box idx = TimedeltaIndex(np.arange(5, dtype="int64")) idx = tm.box_expected(idx, box) @@ -1604,13 +1616,16 @@ def test_td64arr_div_nat_invalid(self, box_with_array): def test_td64arr_div_td64nat(self, box_with_array): # GH#23829 + box = box_with_array + xbox = np.ndarray if box is pd.array else box + rng = timedelta_range("1 days", "10 days") - rng = tm.box_expected(rng, box_with_array) + rng = tm.box_expected(rng, box) other = np.timedelta64("NaT") expected = np.array([np.nan] * 10) - expected = tm.box_expected(expected, box_with_array) + expected = tm.box_expected(expected, xbox) result = rng / other tm.assert_equal(result, expected) @@ -1631,11 +1646,14 @@ def test_td64arr_div_int(self, box_with_array): def test_td64arr_div_tdlike_scalar(self, two_hours, box_with_array): # GH#20088, GH#22163 ensure DataFrame returns correct dtype + box = box_with_array + xbox = np.ndarray if box is pd.array else box + rng = timedelta_range("1 days", "10 days", name="foo") expected = pd.Float64Index((np.arange(10) + 1) * 12, name="foo") - rng = tm.box_expected(rng, box_with_array) - expected = tm.box_expected(expected, box_with_array) + rng = tm.box_expected(rng, box) + expected = tm.box_expected(expected, xbox) result = rng / two_hours tm.assert_equal(result, expected) @@ -1647,32 +1665,38 @@ def test_td64arr_div_tdlike_scalar(self, two_hours, box_with_array): @pytest.mark.parametrize("m", [1, 3, 10]) @pytest.mark.parametrize("unit", ["D", "h", "m", "s", "ms", "us", "ns"]) def test_td64arr_div_td64_scalar(self, m, unit, box_with_array): + box = box_with_array + xbox = np.ndarray if box is pd.array else box + startdate = Series(pd.date_range("2013-01-01", "2013-01-03")) enddate = Series(pd.date_range("2013-03-01", "2013-03-03")) ser = enddate - startdate ser[2] = np.nan flat = ser - ser = tm.box_expected(ser, box_with_array) + ser = tm.box_expected(ser, box) # op expected = Series([x / np.timedelta64(m, unit) for x in flat]) - expected = tm.box_expected(expected, box_with_array) + expected = tm.box_expected(expected, xbox) result = ser / np.timedelta64(m, unit) tm.assert_equal(result, expected) # reverse op expected = Series([Timedelta(np.timedelta64(m, unit)) / x for x in flat]) - expected = tm.box_expected(expected, box_with_array) + expected = tm.box_expected(expected, xbox) result = np.timedelta64(m, unit) / ser tm.assert_equal(result, expected) def test_td64arr_div_tdlike_scalar_with_nat(self, two_hours, box_with_array): + box = box_with_array + xbox = np.ndarray if box is pd.array else box + rng = TimedeltaIndex(["1 days", pd.NaT, "2 days"], name="foo") expected = pd.Float64Index([12, np.nan, 24], name="foo") - rng = tm.box_expected(rng, box_with_array) - expected = tm.box_expected(expected, box_with_array) + rng = tm.box_expected(rng, box) + expected = tm.box_expected(expected, xbox) result = rng / two_hours tm.assert_equal(result, expected) @@ -1683,17 +1707,20 @@ def test_td64arr_div_tdlike_scalar_with_nat(self, two_hours, box_with_array): def test_td64arr_div_td64_ndarray(self, box_with_array): # GH#22631 + box = box_with_array + xbox = np.ndarray if box is pd.array else box + rng = TimedeltaIndex(["1 days", pd.NaT, "2 days"]) expected = pd.Float64Index([12, np.nan, 24]) - rng = tm.box_expected(rng, box_with_array) - expected = tm.box_expected(expected, box_with_array) + rng = tm.box_expected(rng, box) + expected = tm.box_expected(expected, xbox) other = np.array([2, 4, 2], dtype="m8[h]") result = rng / other tm.assert_equal(result, expected) - result = rng / tm.box_expected(other, box_with_array) + result = rng / tm.box_expected(other, box) tm.assert_equal(result, expected) result = rng / other.astype(object) @@ -1707,7 +1734,7 @@ def test_td64arr_div_td64_ndarray(self, box_with_array): result = other / rng tm.assert_equal(result, expected) - result = tm.box_expected(other, box_with_array) / rng + result = tm.box_expected(other, box) / rng tm.assert_equal(result, expected) result = other.astype(object) / rng @@ -1736,6 +1763,7 @@ def test_tdarr_div_length_mismatch(self, box_with_array): def test_td64arr_floordiv_td64arr_with_nat(self, box_with_array): # GH#35529 box = box_with_array + xbox = np.ndarray if box is pd.array else box left = pd.Series([1000, 222330, 30], dtype="timedelta64[ns]") right = pd.Series([1000, 222330, None], dtype="timedelta64[ns]") @@ -1744,7 +1772,7 @@ def test_td64arr_floordiv_td64arr_with_nat(self, box_with_array): right = tm.box_expected(right, box) expected = np.array([1.0, 1.0, np.nan], dtype=np.float64) - expected = tm.box_expected(expected, box) + expected = tm.box_expected(expected, xbox) result = left // right @@ -1756,39 +1784,48 @@ def test_td64arr_floordiv_td64arr_with_nat(self, box_with_array): def test_td64arr_floordiv_tdscalar(self, box_with_array, scalar_td): # GH#18831 + box = box_with_array + xbox = np.ndarray if box is pd.array else box + td1 = Series([timedelta(minutes=5, seconds=3)] * 3) td1.iloc[2] = np.nan expected = Series([0, 0, np.nan]) - td1 = tm.box_expected(td1, box_with_array, transpose=False) - expected = tm.box_expected(expected, box_with_array, transpose=False) + td1 = tm.box_expected(td1, box, transpose=False) + expected = tm.box_expected(expected, xbox, transpose=False) result = td1 // scalar_td tm.assert_equal(result, expected) def test_td64arr_rfloordiv_tdscalar(self, box_with_array, scalar_td): # GH#18831 + box = box_with_array + xbox = np.ndarray if box is pd.array else box + td1 = Series([timedelta(minutes=5, seconds=3)] * 3) td1.iloc[2] = np.nan expected = Series([1, 1, np.nan]) - td1 = tm.box_expected(td1, box_with_array, transpose=False) - expected = tm.box_expected(expected, box_with_array, transpose=False) + td1 = tm.box_expected(td1, box, transpose=False) + expected = tm.box_expected(expected, xbox, transpose=False) result = scalar_td // td1 tm.assert_equal(result, expected) def test_td64arr_rfloordiv_tdscalar_explicit(self, box_with_array, scalar_td): # GH#18831 + box = box_with_array + xbox = np.ndarray if box is pd.array else box + td1 = Series([timedelta(minutes=5, seconds=3)] * 3) td1.iloc[2] = np.nan expected = Series([1, 1, np.nan]) - td1 = tm.box_expected(td1, box_with_array, transpose=False) - expected = tm.box_expected(expected, box_with_array, transpose=False) + td1 = tm.box_expected(td1, box, transpose=False) + expected = tm.box_expected(expected, xbox, transpose=False) # We can test __rfloordiv__ using this syntax, # see `test_timedelta_rfloordiv` @@ -1806,11 +1843,14 @@ def test_td64arr_floordiv_int(self, box_with_array): 1 // idx def test_td64arr_floordiv_tdlike_scalar(self, two_hours, box_with_array): + box = box_with_array + xbox = np.ndarray if box is pd.array else box + tdi = timedelta_range("1 days", "10 days", name="foo") expected = pd.Int64Index((np.arange(10) + 1) * 12, name="foo") - tdi = tm.box_expected(tdi, box_with_array) - expected = tm.box_expected(expected, box_with_array) + tdi = tm.box_expected(tdi, box) + expected = tm.box_expected(expected, xbox) result = tdi // two_hours tm.assert_equal(result, expected) @@ -1827,17 +1867,20 @@ def test_td64arr_floordiv_tdlike_scalar(self, two_hours, box_with_array): ) def test_td64arr_rfloordiv_tdlike_scalar(self, scalar_td, box_with_array): # GH#19125 + box = box_with_array + xbox = np.ndarray if box_with_array is pd.array else box_with_array + tdi = TimedeltaIndex(["00:05:03", "00:05:03", pd.NaT], freq=None) expected = pd.Index([2.0, 2.0, np.nan]) - tdi = tm.box_expected(tdi, box_with_array, transpose=False) - expected = tm.box_expected(expected, box_with_array, transpose=False) + tdi = tm.box_expected(tdi, box, transpose=False) + expected = tm.box_expected(expected, xbox, transpose=False) res = tdi.__rfloordiv__(scalar_td) tm.assert_equal(res, expected) expected = pd.Index([0.0, 0.0, np.nan]) - expected = tm.box_expected(expected, box_with_array, transpose=False) + expected = tm.box_expected(expected, xbox, transpose=False) res = tdi // (scalar_td) tm.assert_equal(res, expected) @@ -2059,7 +2102,7 @@ def test_td64arr_mul_int_series(self, box_with_array, names, request): reason = "broadcasts along wrong axis, but doesn't raise" request.node.add_marker(pytest.mark.xfail(reason=reason)) - exname = names[2] if box is not tm.to_array else names[1] + exname = names[2] if box not in [tm.to_array, pd.array] else names[1] tdi = TimedeltaIndex( ["0days", "1day", "2days", "3days", "4days"], name=names[0] @@ -2074,8 +2117,12 @@ def test_td64arr_mul_int_series(self, box_with_array, names, request): ) tdi = tm.box_expected(tdi, box) - box = Series if (box is pd.Index or box is tm.to_array) else box - expected = tm.box_expected(expected, box) + xbox = ( + Series + if (box is pd.Index or box is tm.to_array or box is pd.array) + else box + ) + expected = tm.box_expected(expected, xbox) result = ser * tdi tm.assert_equal(result, expected) @@ -2098,7 +2145,7 @@ def test_float_series_rdiv_td64arr(self, box_with_array, names): ) ser = Series([1.5, 3, 4.5, 6, 7.5], dtype=np.float64, name=names[1]) - xname = names[2] if box is not tm.to_array else names[1] + xname = names[2] if box not in [tm.to_array, pd.array] else names[1] expected = Series( [tdi[n] / ser[n] for n in range(len(ser))], dtype="timedelta64[ns]", @@ -2106,7 +2153,7 @@ def test_float_series_rdiv_td64arr(self, box_with_array, names): ) xbox = box - if box in [pd.Index, tm.to_array] and type(ser) is Series: + if box in [pd.Index, tm.to_array, pd.array] and type(ser) is Series: xbox = Series tdi = tm.box_expected(tdi, box)