TST: parametrize and de-duplicate arith tests (pandas-dev#27950)

Author: jbrockmendel

pandas/tests/arithmetic/conftest.py

@@ -190,7 +190,12 @@ def box(request):
 
 
 @pytest.fixture(
-    params=[pd.Index, pd.Series, pytest.param(pd.DataFrame, marks=pytest.mark.xfail)],
+    params=[
+        pd.Index,
+        pd.Series,
+        pytest.param(pd.DataFrame, marks=pytest.mark.xfail),
+        tm.to_array,
+    ],
     ids=id_func,
 )
 def box_df_fail(request):
@@ -206,6 +211,7 @@ def box_df_fail(request):
         (pd.Series, False),
         (pd.DataFrame, False),
         pytest.param((pd.DataFrame, True), marks=pytest.mark.xfail),
+        (tm.to_array, False),
     ],
     ids=id_func,
 )

pandas/tests/arithmetic/test_datetime64.py

@@ -348,28 +348,6 @@ def test_dt64arr_timestamp_equality(self, box_with_array):
         expected = tm.box_expected([False, False], xbox)
         tm.assert_equal(result, expected)
 
-    @pytest.mark.parametrize(
-        "op",
-        [operator.eq, operator.ne, operator.gt, operator.ge, operator.lt, operator.le],
-    )
-    def test_comparison_tzawareness_compat(self, op):
-        # GH#18162
-        dr = pd.date_range("2016-01-01", periods=6)
-        dz = dr.tz_localize("US/Pacific")
-
-        # Check that there isn't a problem aware-aware and naive-naive do not
-        # raise
-        naive_series = Series(dr)
-        aware_series = Series(dz)
-        msg = "Cannot compare tz-naive and tz-aware"
-        with pytest.raises(TypeError, match=msg):
-            op(dz, naive_series)
-        with pytest.raises(TypeError, match=msg):
-            op(dr, aware_series)
-
-        # TODO: implement _assert_tzawareness_compat for the reverse
-        # comparison with the Series on the left-hand side
-
 
 class TestDatetimeIndexComparisons:
 
@@ -599,15 +577,18 @@ def test_comparison_tzawareness_compat(self, op, box_df_fail):
         with pytest.raises(TypeError, match=msg):
             op(dz, np.array(list(dr), dtype=object))
 
-        # Check that there isn't a problem aware-aware and naive-naive do not
-        # raise
+        # The aware==aware and naive==naive comparisons should *not* raise
         assert_all(dr == dr)
-        assert_all(dz == dz)
+        assert_all(dr == list(dr))
+        assert_all(list(dr) == dr)
+        assert_all(np.array(list(dr), dtype=object) == dr)
+        assert_all(dr == np.array(list(dr), dtype=object))
 
-        # FIXME: DataFrame case fails to raise for == and !=, wrong
-        #  message for inequalities
-        assert (dr == list(dr)).all()
-        assert (dz == list(dz)).all()
+        assert_all(dz == dz)
+        assert_all(dz == list(dz))
+        assert_all(list(dz) == dz)
+        assert_all(np.array(list(dz), dtype=object) == dz)
+        assert_all(dz == np.array(list(dz), dtype=object))
 
     @pytest.mark.parametrize(
         "op",
@@ -844,6 +825,7 @@ def test_dt64arr_isub_timedeltalike_scalar(
         rng -= two_hours
         tm.assert_equal(rng, expected)
 
+    # TODO: redundant with test_dt64arr_add_timedeltalike_scalar
     def test_dt64arr_add_td64_scalar(self, box_with_array):
         # scalar timedeltas/np.timedelta64 objects
         # operate with np.timedelta64 correctly
@@ -1709,14 +1691,12 @@ def test_operators_datetimelike(self):
         dt1 - dt2
         dt2 - dt1
 
-        # ## datetime64 with timetimedelta ###
+        # datetime64 with timetimedelta
         dt1 + td1
         td1 + dt1
         dt1 - td1
-        # TODO: Decide if this ought to work.
-        # td1 - dt1
 
-        # ## timetimedelta with datetime64 ###
+        # timetimedelta with datetime64
         td1 + dt1
         dt1 + td1
 
@@ -1914,7 +1894,7 @@ def test_dt64_series_add_intlike(self, tz, op):
         with pytest.raises(TypeError, match=msg):
             method(other)
         with pytest.raises(TypeError, match=msg):
-            method(other.values)
+            method(np.array(other))
         with pytest.raises(TypeError, match=msg):
             method(pd.Index(other))
 
@@ -2380,34 +2360,34 @@ def test_ufunc_coercions(self):
         idx = date_range("2011-01-01", periods=3, freq="2D", name="x")
 
         delta = np.timedelta64(1, "D")
+        exp = date_range("2011-01-02", periods=3, freq="2D", name="x")
         for result in [idx + delta, np.add(idx, delta)]:
             assert isinstance(result, DatetimeIndex)
-            exp = date_range("2011-01-02", periods=3, freq="2D", name="x")
             tm.assert_index_equal(result, exp)
             assert result.freq == "2D"
 
+        exp = date_range("2010-12-31", periods=3, freq="2D", name="x")
         for result in [idx - delta, np.subtract(idx, delta)]:
             assert isinstance(result, DatetimeIndex)
-            exp = date_range("2010-12-31", periods=3, freq="2D", name="x")
             tm.assert_index_equal(result, exp)
             assert result.freq == "2D"
 
         delta = np.array(
             [np.timedelta64(1, "D"), np.timedelta64(2, "D"), np.timedelta64(3, "D")]
         )
+        exp = DatetimeIndex(
+            ["2011-01-02", "2011-01-05", "2011-01-08"], freq="3D", name="x"
+        )
         for result in [idx + delta, np.add(idx, delta)]:
             assert isinstance(result, DatetimeIndex)
-            exp = DatetimeIndex(
-                ["2011-01-02", "2011-01-05", "2011-01-08"], freq="3D", name="x"
-            )
             tm.assert_index_equal(result, exp)
             assert result.freq == "3D"
 
+        exp = DatetimeIndex(
+            ["2010-12-31", "2011-01-01", "2011-01-02"], freq="D", name="x"
+        )
         for result in [idx - delta, np.subtract(idx, delta)]:
             assert isinstance(result, DatetimeIndex)
-            exp = DatetimeIndex(
-                ["2010-12-31", "2011-01-01", "2011-01-02"], freq="D", name="x"
-            )
             tm.assert_index_equal(result, exp)
             assert result.freq == "D"
 

pandas/tests/arithmetic/test_numeric.py

@@ -561,9 +561,9 @@ def test_div_int(self, numeric_idx):
         tm.assert_index_equal(result, expected)
 
     @pytest.mark.parametrize("op", [operator.mul, ops.rmul, operator.floordiv])
-    def test_mul_int_identity(self, op, numeric_idx, box):
+    def test_mul_int_identity(self, op, numeric_idx, box_with_array):
         idx = numeric_idx
-        idx = tm.box_expected(idx, box)
+        idx = tm.box_expected(idx, box_with_array)
 
         result = op(idx, 1)
         tm.assert_equal(result, idx)
@@ -615,8 +615,9 @@ def test_mul_size_mismatch_raises(self, numeric_idx):
             idx * np.array([1, 2])
 
     @pytest.mark.parametrize("op", [operator.pow, ops.rpow])
-    def test_pow_float(self, op, numeric_idx, box):
+    def test_pow_float(self, op, numeric_idx, box_with_array):
         # test power calculations both ways, GH#14973
+        box = box_with_array
         idx = numeric_idx
         expected = pd.Float64Index(op(idx.values, 2.0))
 
@@ -626,8 +627,9 @@ def test_pow_float(self, op, numeric_idx, box):
         result = op(idx, 2.0)
         tm.assert_equal(result, expected)
 
-    def test_modulo(self, numeric_idx, box):
+    def test_modulo(self, numeric_idx, box_with_array):
         # GH#9244
+        box = box_with_array
         idx = numeric_idx
         expected = Index(idx.values % 2)
 
@@ -1041,7 +1043,8 @@ class TestObjectDtypeEquivalence:
     # Tests that arithmetic operations match operations executed elementwise
 
     @pytest.mark.parametrize("dtype", [None, object])
-    def test_numarr_with_dtype_add_nan(self, dtype, box):
+    def test_numarr_with_dtype_add_nan(self, dtype, box_with_array):
+        box = box_with_array
         ser = pd.Series([1, 2, 3], dtype=dtype)
         expected = pd.Series([np.nan, np.nan, np.nan], dtype=dtype)
 
@@ -1055,7 +1058,8 @@ def test_numarr_with_dtype_add_nan(self, dtype, box):
         tm.assert_equal(result, expected)
 
     @pytest.mark.parametrize("dtype", [None, object])
-    def test_numarr_with_dtype_add_int(self, dtype, box):
+    def test_numarr_with_dtype_add_int(self, dtype, box_with_array):
+        box = box_with_array
         ser = pd.Series([1, 2, 3], dtype=dtype)
         expected = pd.Series([2, 3, 4], dtype=dtype)
 

pandas/tests/arithmetic/test_object.py

@@ -89,16 +89,16 @@ def test_pow_ops_object(self):
 
     @pytest.mark.parametrize("op", [operator.add, ops.radd])
     @pytest.mark.parametrize("other", ["category", "Int64"])
-    def test_add_extension_scalar(self, other, box, op):
+    def test_add_extension_scalar(self, other, box_with_array, op):
         # GH#22378
         # Check that scalars satisfying is_extension_array_dtype(obj)
         # do not incorrectly try to dispatch to an ExtensionArray operation
 
         arr = pd.Series(["a", "b", "c"])
         expected = pd.Series([op(x, other) for x in arr])
 
-        arr = tm.box_expected(arr, box)
-        expected = tm.box_expected(expected, box)
+        arr = tm.box_expected(arr, box_with_array)
+        expected = tm.box_expected(expected, box_with_array)
 
         result = op(arr, other)
         tm.assert_equal(result, expected)
@@ -133,16 +133,17 @@ def test_objarr_radd_str(self, box):
         ],
     )
     @pytest.mark.parametrize("dtype", [None, object])
-    def test_objarr_radd_str_invalid(self, dtype, data, box):
+    def test_objarr_radd_str_invalid(self, dtype, data, box_with_array):
         ser = Series(data, dtype=dtype)
 
-        ser = tm.box_expected(ser, box)
+        ser = tm.box_expected(ser, box_with_array)
         with pytest.raises(TypeError):
             "foo_" + ser
 
     @pytest.mark.parametrize("op", [operator.add, ops.radd, operator.sub, ops.rsub])
-    def test_objarr_add_invalid(self, op, box):
+    def test_objarr_add_invalid(self, op, box_with_array):
         # invalid ops
+        box = box_with_array
 
         obj_ser = tm.makeObjectSeries()
         obj_ser.name = "objects"