TST: make get_upcast_box more flexible (#44385)

Author: jbrockmendel

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)

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):

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)

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)

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,14 +2000,15 @@ 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)
 
         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,14 +2027,14 @@ 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)
 
         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)