TST: collect arithmetic test helpers (#30354)

Author: jbrockmendel

pandas/tests/arithmetic/common.py

@@ -0,0 +1,89 @@
+"""
+Assertion helpers for arithmetic tests.
+"""
+import numpy as np
+import pytest
+
+from pandas import DataFrame, Index, Series
+import pandas.util.testing as tm
+
+
+def assert_invalid_addsub_type(left, right, msg=None):
+    """
+    Helper to assert that left and right can be neither added nor subtracted.
+
+    Parameters
+    ---------
+    left : object
+    right : object
+    msg : str or None, default None
+    """
+    with pytest.raises(TypeError, match=msg):
+        left + right
+    with pytest.raises(TypeError, match=msg):
+        right + left
+    with pytest.raises(TypeError, match=msg):
+        left - right
+    with pytest.raises(TypeError, match=msg):
+        right - left
+
+
+def get_upcast_box(box, vector):
+    """
+    Given two box-types, find the one that takes priority
+    """
+    if box is DataFrame or isinstance(vector, DataFrame):
+        return DataFrame
+    if box is Series or isinstance(vector, Series):
+        return Series
+    if box is Index or isinstance(vector, Index):
+        return Index
+    return box
+
+
+def assert_invalid_comparison(left, right, box):
+    """
+    Assert that comparison operations with mismatched types behave correctly.
+
+    Parameters
+    ----------
+    left : np.ndarray, ExtensionArray, Index, or Series
+    right : object
+    box : {pd.DataFrame, pd.Series, pd.Index, 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
+
+    result = left == right
+    expected = xbox(np.zeros(result.shape, dtype=np.bool_))
+
+    tm.assert_equal(result, expected)
+
+    result = right == left
+    tm.assert_equal(result, expected)
+
+    result = left != right
+    tm.assert_equal(result, ~expected)
+
+    result = right != left
+    tm.assert_equal(result, ~expected)
+
+    msg = "Invalid comparison between"
+    with pytest.raises(TypeError, match=msg):
+        left < right
+    with pytest.raises(TypeError, match=msg):
+        left <= right
+    with pytest.raises(TypeError, match=msg):
+        left > right
+    with pytest.raises(TypeError, match=msg):
+        left >= right
+    with pytest.raises(TypeError, match=msg):
+        right < left
+    with pytest.raises(TypeError, match=msg):
+        right <= left
+    with pytest.raises(TypeError, match=msg):
+        right > left
+    with pytest.raises(TypeError, match=msg):
+        right >= left

pandas/tests/arithmetic/test_datetime64.py

@@ -29,57 +29,13 @@
 import pandas.core.arrays.datetimelike as dtl
 from pandas.core.indexes.datetimes import _to_M8
 from pandas.core.ops import roperator
+from pandas.tests.arithmetic.common import (
+    assert_invalid_addsub_type,
+    assert_invalid_comparison,
+    get_upcast_box,
+)
 import pandas.util.testing as tm
 
-
-def assert_invalid_comparison(left, right, box):
-    """
-    Assert that comparison operations with mismatched types behave correctly.
-
-    Parameters
-    ----------
-    left : np.ndarray, ExtensionArray, Index, or Series
-    right : object
-    box : {pd.DataFrame, pd.Series, pd.Index, 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 pd.Index else np.array
-
-    result = left == right
-    expected = xbox(np.zeros(result.shape, dtype=np.bool_))
-
-    tm.assert_equal(result, expected)
-
-    result = right == left
-    tm.assert_equal(result, expected)
-
-    result = left != right
-    tm.assert_equal(result, ~expected)
-
-    result = right != left
-    tm.assert_equal(result, ~expected)
-
-    msg = "Invalid comparison between"
-    with pytest.raises(TypeError, match=msg):
-        left < right
-    with pytest.raises(TypeError, match=msg):
-        left <= right
-    with pytest.raises(TypeError, match=msg):
-        left > right
-    with pytest.raises(TypeError, match=msg):
-        left >= right
-    with pytest.raises(TypeError, match=msg):
-        right < left
-    with pytest.raises(TypeError, match=msg):
-        right <= left
-    with pytest.raises(TypeError, match=msg):
-        right > left
-    with pytest.raises(TypeError, match=msg):
-        right >= left
-
-
 # ------------------------------------------------------------------
 # Comparisons
 
@@ -1033,14 +989,7 @@ def test_dt64arr_add_sub_invalid(self, dti_freq, other, box_with_array):
                 "ufunc '?(add|subtract)'? cannot use operands with types",
             ]
         )
-        with pytest.raises(TypeError, match=msg):
-            dtarr + other
-        with pytest.raises(TypeError, match=msg):
-            other + dtarr
-        with pytest.raises(TypeError, match=msg):
-            dtarr - other
-        with pytest.raises(TypeError, match=msg):
-            other - dtarr
+        assert_invalid_addsub_type(dtarr, other, msg)
 
     @pytest.mark.parametrize("pi_freq", ["D", "W", "Q", "H"])
     @pytest.mark.parametrize("dti_freq", [None, "D"])
@@ -1061,14 +1010,7 @@ def test_dt64arr_add_sub_parr(
                 "ufunc.*cannot use operands",
             ]
         )
-        with pytest.raises(TypeError, match=msg):
-            dtarr + parr
-        with pytest.raises(TypeError, match=msg):
-            parr + dtarr
-        with pytest.raises(TypeError, match=msg):
-            dtarr - parr
-        with pytest.raises(TypeError, match=msg):
-            parr - dtarr
+        assert_invalid_addsub_type(dtarr, parr, msg)
 
 
 class TestDatetime64DateOffsetArithmetic:
@@ -2368,7 +2310,6 @@ def test_dti_addsub_offset_arraylike(
         # GH#18849, GH#19744
         box = pd.Index
         other_box = index_or_series
-        from .test_timedelta64 import get_upcast_box
 
         tz = tz_naive_fixture
         dti = pd.date_range("2017-01-01", periods=2, tz=tz, name=names[0])

pandas/tests/arithmetic/test_timedelta64.py

@@ -18,23 +18,9 @@
     Timestamp,
     timedelta_range,
 )
-from pandas.tests.arithmetic.test_datetime64 import assert_invalid_comparison
+from pandas.tests.arithmetic.common import assert_invalid_comparison, get_upcast_box
 import pandas.util.testing as tm
 
-
-def get_upcast_box(box, vector):
-    """
-    Given two box-types, find the one that takes priority
-    """
-    if box is DataFrame or isinstance(vector, DataFrame):
-        return DataFrame
-    if box is Series or isinstance(vector, Series):
-        return Series
-    if box is pd.Index or isinstance(vector, pd.Index):
-        return pd.Index
-    return box
-
-
 # ------------------------------------------------------------------
 # Timedelta64[ns] dtype Comparisons