diff --git a/pandas/tests/arithmetic/common.py b/pandas/tests/arithmetic/common.py new file mode 100644 index 0000000000000..bc02a1e76a695 --- /dev/null +++ b/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 diff --git a/pandas/tests/arithmetic/test_datetime64.py b/pandas/tests/arithmetic/test_datetime64.py index b77c9a2bddcfa..c055b3e62a368 100644 --- a/pandas/tests/arithmetic/test_datetime64.py +++ b/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]) diff --git a/pandas/tests/arithmetic/test_timedelta64.py b/pandas/tests/arithmetic/test_timedelta64.py index 4a37a56f5029c..cafbbc9aef6f7 100644 --- a/pandas/tests/arithmetic/test_timedelta64.py +++ b/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