Skip to content

TST: collect arithmetic test helpers #30354

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Dec 23, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 89 additions & 0 deletions pandas/tests/arithmetic/common.py
Original file line number Diff line number Diff line change
@@ -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
73 changes: 7 additions & 66 deletions pandas/tests/arithmetic/test_datetime64.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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"])
Expand All @@ -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:
Expand Down Expand Up @@ -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])
Expand Down
16 changes: 1 addition & 15 deletions pandas/tests/arithmetic/test_timedelta64.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down