From 982a787117f1d67c899e13d28b871aa142f40630 Mon Sep 17 00:00:00 2001 From: Brock Date: Wed, 10 Nov 2021 12:22:40 -0800 Subject: [PATCH] TST: make get_upcast_box more flexible --- pandas/_testing/__init__.py | 2 +- pandas/tests/arithmetic/common.py | 31 +++++++++++---------- pandas/tests/arithmetic/test_datetime64.py | 21 ++++++-------- pandas/tests/arithmetic/test_period.py | 19 +++++++------ pandas/tests/arithmetic/test_timedelta64.py | 16 +++++------ 5 files changed, 44 insertions(+), 45 deletions(-) diff --git a/pandas/_testing/__init__.py b/pandas/_testing/__init__.py index e8283a222d86a..c2c55a4060f7a 100644 --- a/pandas/_testing/__init__.py +++ b/pandas/_testing/__init__.py @@ -259,7 +259,7 @@ def box_expected(expected, box_cls, transpose=True): expected = DatetimeArray(expected) elif box_cls is TimedeltaArray: expected = TimedeltaArray(expected) - elif box_cls is np.ndarray: + elif box_cls is np.ndarray or box_cls is np.array: expected = np.array(expected) elif box_cls is to_array: expected = to_array(expected) diff --git a/pandas/tests/arithmetic/common.py b/pandas/tests/arithmetic/common.py index 6f4e35ad4dfb2..af70cdfe538bb 100644 --- a/pandas/tests/arithmetic/common.py +++ b/pandas/tests/arithmetic/common.py @@ -34,26 +34,29 @@ def assert_invalid_addsub_type(left, right, msg=None): right - left -def get_expected_box(box): +def get_upcast_box(left, right, is_cmp: bool = False): """ - Get the box to use for 'expected' in a comparison operation. - """ - if box in [Index, array]: - return np.ndarray - return box - + Get the box to use for 'expected' in an arithmetic or comparison operation. -def get_upcast_box(box, vector): - """ - Given two box-types, find the one that takes priority. + Parameters + left : Any + right : Any + is_cmp : bool, default False + Whether the operation is a comparison method. """ - if box is DataFrame or isinstance(vector, DataFrame): + + if isinstance(left, DataFrame) or isinstance(right, DataFrame): return DataFrame - if box is Series or isinstance(vector, Series): + if isinstance(left, Series) or isinstance(right, Series): + if is_cmp and isinstance(left, Index): + # Index does not defer for comparisons + return np.array return Series - if box is Index or isinstance(vector, Index): + if isinstance(left, Index) or isinstance(right, Index): + if is_cmp: + return np.array return Index - return box + return tm.to_array def assert_invalid_comparison(left, right, box): diff --git a/pandas/tests/arithmetic/test_datetime64.py b/pandas/tests/arithmetic/test_datetime64.py index 82f1e60f0aea5..44a70d3933b66 100644 --- a/pandas/tests/arithmetic/test_datetime64.py +++ b/pandas/tests/arithmetic/test_datetime64.py @@ -43,7 +43,6 @@ from pandas.tests.arithmetic.common import ( assert_invalid_addsub_type, assert_invalid_comparison, - get_expected_box, get_upcast_box, ) @@ -60,12 +59,12 @@ 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 = get_expected_box(box) dti = date_range("20130101", periods=3, tz=tz) other = np.array(dti.to_numpy()[0]) dtarr = tm.box_expected(dti, box) + xbox = get_upcast_box(dtarr, other, True) result = dtarr <= other expected = np.array([True, False, False]) expected = tm.box_expected(expected, xbox) @@ -147,12 +146,12 @@ 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 = get_expected_box(box) ts = Timestamp.now(tz) ser = Series([ts, NaT]) obj = tm.box_expected(ser, box) + xbox = get_upcast_box(obj, ts, True) expected = Series([True, False], dtype=np.bool_) expected = tm.box_expected(expected, xbox) @@ -244,10 +243,9 @@ def test_nat_comparisons_scalar(self, dtype, data, box_with_array): # on older numpys (since they check object identity) return - xbox = get_expected_box(box) - left = Series(data, dtype=dtype) left = tm.box_expected(left, box) + xbox = get_upcast_box(left, NaT, True) expected = [False, False, False] expected = tm.box_expected(expected, xbox) @@ -323,10 +321,10 @@ def test_timestamp_compare_series(self, left, right): def test_dt64arr_timestamp_equality(self, box_with_array): # GH#11034 - xbox = get_expected_box(box_with_array) ser = Series([Timestamp("2000-01-29 01:59:00"), Timestamp("2000-01-30"), NaT]) ser = tm.box_expected(ser, box_with_array) + xbox = get_upcast_box(ser, ser, True) result = ser != ser expected = tm.box_expected([False, False, True], xbox) @@ -417,13 +415,12 @@ def test_dti_cmp_nat(self, dtype, box_with_array): # on older numpys (since they check object identity) return - xbox = get_expected_box(box_with_array) - left = DatetimeIndex([Timestamp("2011-01-01"), NaT, Timestamp("2011-01-03")]) right = DatetimeIndex([NaT, NaT, Timestamp("2011-01-03")]) left = tm.box_expected(left, box_with_array) right = tm.box_expected(right, box_with_array) + xbox = get_upcast_box(left, right, True) lhs, rhs = left, right if dtype is object: @@ -642,12 +639,11 @@ def test_scalar_comparison_tzawareness( self, comparison_op, other, tz_aware_fixture, box_with_array ): op = comparison_op - box = box_with_array tz = tz_aware_fixture dti = date_range("2016-01-01", periods=2, tz=tz) - xbox = get_expected_box(box) dtarr = tm.box_expected(dti, box_with_array) + xbox = get_upcast_box(dtarr, other, True) if op in [operator.eq, operator.ne]: exbool = op is operator.ne expected = np.array([exbool, exbool], dtype=bool) @@ -2421,14 +2417,13 @@ def test_dti_addsub_offset_arraylike( self, tz_naive_fixture, names, op, index_or_series ): # GH#18849, GH#19744 - box = pd.Index other_box = index_or_series tz = tz_naive_fixture dti = date_range("2017-01-01", periods=2, tz=tz, name=names[0]) other = other_box([pd.offsets.MonthEnd(), pd.offsets.Day(n=2)], name=names[1]) - xbox = get_upcast_box(box, other) + xbox = get_upcast_box(dti, other) with tm.assert_produces_warning(PerformanceWarning): res = op(dti, other) @@ -2448,7 +2443,7 @@ def test_dti_addsub_object_arraylike( dti = date_range("2017-01-01", periods=2, tz=tz) dtarr = tm.box_expected(dti, box_with_array) other = other_box([pd.offsets.MonthEnd(), Timedelta(days=4)]) - xbox = get_upcast_box(box_with_array, other) + xbox = get_upcast_box(dtarr, other) expected = DatetimeIndex(["2017-01-31", "2017-01-06"], tz=tz_naive_fixture) expected = tm.box_expected(expected, xbox) diff --git a/pandas/tests/arithmetic/test_period.py b/pandas/tests/arithmetic/test_period.py index 41c2cb2cc4f1e..f8814a33292ec 100644 --- a/pandas/tests/arithmetic/test_period.py +++ b/pandas/tests/arithmetic/test_period.py @@ -27,7 +27,7 @@ from pandas.core.arrays import TimedeltaArray from pandas.tests.arithmetic.common import ( assert_invalid_comparison, - get_expected_box, + get_upcast_box, ) # ------------------------------------------------------------------ @@ -41,12 +41,13 @@ class TestPeriodArrayLikeComparisons: def test_compare_zerodim(self, box_with_array): # GH#26689 make sure we unbox zero-dimensional arrays - xbox = get_expected_box(box_with_array) pi = period_range("2000", periods=4) other = np.array(pi.to_numpy()[0]) pi = tm.box_expected(pi, box_with_array) + xbox = get_upcast_box(pi, other, True) + result = pi <= other expected = np.array([True, False, False, False]) expected = tm.box_expected(expected, xbox) @@ -78,11 +79,11 @@ def test_compare_invalid_listlike(self, box_with_array, other): @pytest.mark.parametrize("other_box", [list, np.array, lambda x: x.astype(object)]) def test_compare_object_dtype(self, box_with_array, other_box): - xbox = get_expected_box(box_with_array) pi = period_range("2000", periods=5) parr = tm.box_expected(pi, box_with_array) other = other_box(pi) + xbox = get_upcast_box(parr, other, True) expected = np.array([True, True, True, True, True]) expected = tm.box_expected(expected, xbox) @@ -195,14 +196,15 @@ 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 = get_expected_box(box_with_array) - pi = period_range("2000-01-01", periods=10, freq="D") val = Period("2000-01-04", freq="D") + expected = [x > val for x in pi] ser = tm.box_expected(pi, box_with_array) + xbox = get_upcast_box(ser, val, True) + expected = tm.box_expected(expected, xbox) result = ser > val tm.assert_equal(result, expected) @@ -216,11 +218,10 @@ 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 = get_expected_box(box_with_array) - base = PeriodIndex(["2011-01", "2011-02", "2011-03", "2011-04"], freq=freq) base = tm.box_expected(base, box_with_array) per = Period("2011-02", freq=freq) + xbox = get_upcast_box(base, per, True) exp = np.array([False, True, False, False]) exp = tm.box_expected(exp, xbox) @@ -255,14 +256,14 @@ 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 = get_expected_box(box_with_array) - base = PeriodIndex(["2011-01", "2011-02", "2011-03", "2011-04"], freq=freq) base = tm.box_expected(base, box_with_array) # TODO: could also box idx? idx = PeriodIndex(["2011-02", "2011-01", "2011-03", "2011-05"], freq=freq) + xbox = get_upcast_box(base, idx, True) + exp = np.array([False, False, True, False]) exp = tm.box_expected(exp, xbox) tm.assert_equal(base == idx, exp) diff --git a/pandas/tests/arithmetic/test_timedelta64.py b/pandas/tests/arithmetic/test_timedelta64.py index b8fa6c79b1b93..86980ad42766e 100644 --- a/pandas/tests/arithmetic/test_timedelta64.py +++ b/pandas/tests/arithmetic/test_timedelta64.py @@ -1542,13 +1542,13 @@ def test_tdi_mul_float_series(self, box_with_array): ) def test_tdi_rmul_arraylike(self, other, box_with_array): box = box_with_array - xbox = get_upcast_box(box, other) tdi = TimedeltaIndex(["1 Day"] * 10) - expected = timedelta_range("1 days", "10 days") - expected._data.freq = None + expected = timedelta_range("1 days", "10 days")._with_freq(None) tdi = tm.box_expected(tdi, box) + xbox = get_upcast_box(tdi, other) + expected = tm.box_expected(expected, xbox) result = other * tdi @@ -2000,7 +2000,6 @@ def test_td64arr_rmul_numeric_array( ): # GH#4521 # divide/multiply by integers - xbox = get_upcast_box(box_with_array, vector) tdser = Series(["59 Days", "59 Days", "NaT"], dtype="m8[ns]") vector = vector.astype(any_real_numpy_dtype) @@ -2008,6 +2007,8 @@ def test_td64arr_rmul_numeric_array( expected = Series(["1180 Days", "1770 Days", "NaT"], dtype="timedelta64[ns]") tdser = tm.box_expected(tdser, box_with_array) + xbox = get_upcast_box(tdser, vector) + expected = tm.box_expected(expected, xbox) result = tdser * vector @@ -2026,7 +2027,6 @@ def test_td64arr_div_numeric_array( ): # GH#4521 # divide/multiply by integers - xbox = get_upcast_box(box_with_array, vector) tdser = Series(["59 Days", "59 Days", "NaT"], dtype="m8[ns]") vector = vector.astype(any_real_numpy_dtype) @@ -2034,6 +2034,7 @@ def test_td64arr_div_numeric_array( expected = Series(["2.95D", "1D 23H 12m", "NaT"], dtype="timedelta64[ns]") tdser = tm.box_expected(tdser, box_with_array) + xbox = get_upcast_box(tdser, vector) expected = tm.box_expected(expected, xbox) result = tdser / vector @@ -2085,7 +2086,7 @@ def test_td64arr_mul_int_series(self, box_with_array, names): ) tdi = tm.box_expected(tdi, box) - xbox = get_upcast_box(box, ser) + xbox = get_upcast_box(tdi, ser) expected = tm.box_expected(expected, xbox) @@ -2117,9 +2118,8 @@ def test_float_series_rdiv_td64arr(self, box_with_array, names): name=xname, ) - xbox = get_upcast_box(box, ser) - tdi = tm.box_expected(tdi, box) + xbox = get_upcast_box(tdi, ser) expected = tm.box_expected(expected, xbox) result = ser.__rtruediv__(tdi)