Skip to content

Commit cb31838

Browse files
jbrockmendelvictor
authored and
victor
committed
[POC] implement test_arithmetic.py (pandas-dev#22033)
1 parent f228686 commit cb31838

File tree

4 files changed

+164
-56
lines changed

4 files changed

+164
-56
lines changed

pandas/tests/indexes/test_numeric.py

+1-37
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
import pandas.util.testing as tm
1414

1515
import pandas as pd
16-
from pandas._libs.tslibs import Timestamp, Timedelta
16+
from pandas._libs.tslibs import Timestamp
1717

1818
from pandas.tests.indexes.common import Base
1919

@@ -26,42 +26,6 @@ def full_like(array, value):
2626
return ret
2727

2828

29-
class TestIndexArithmeticWithTimedeltaScalar(object):
30-
31-
@pytest.mark.parametrize('index', [
32-
Int64Index(range(1, 11)),
33-
UInt64Index(range(1, 11)),
34-
Float64Index(range(1, 11)),
35-
RangeIndex(1, 11)])
36-
@pytest.mark.parametrize('scalar_td', [Timedelta(days=1),
37-
Timedelta(days=1).to_timedelta64(),
38-
Timedelta(days=1).to_pytimedelta()])
39-
def test_index_mul_timedelta(self, scalar_td, index):
40-
# GH#19333
41-
expected = pd.timedelta_range('1 days', '10 days')
42-
43-
result = index * scalar_td
44-
tm.assert_index_equal(result, expected)
45-
commute = scalar_td * index
46-
tm.assert_index_equal(commute, expected)
47-
48-
@pytest.mark.parametrize('index', [Int64Index(range(1, 3)),
49-
UInt64Index(range(1, 3)),
50-
Float64Index(range(1, 3)),
51-
RangeIndex(1, 3)])
52-
@pytest.mark.parametrize('scalar_td', [Timedelta(days=1),
53-
Timedelta(days=1).to_timedelta64(),
54-
Timedelta(days=1).to_pytimedelta()])
55-
def test_index_rdiv_timedelta(self, scalar_td, index):
56-
expected = pd.TimedeltaIndex(['1 Day', '12 Hours'])
57-
58-
result = scalar_td / index
59-
tm.assert_index_equal(result, expected)
60-
61-
with pytest.raises(TypeError):
62-
index / scalar_td
63-
64-
6529
class Numeric(Base):
6630

6731
def test_can_hold_identifiers(self):

pandas/tests/series/test_arithmetic.py

-19
Original file line numberDiff line numberDiff line change
@@ -891,22 +891,3 @@ def test_td64series_mul_timedeltalike_invalid(self, scalar_td):
891891
td1 * scalar_td
892892
with tm.assert_raises_regex(TypeError, pattern):
893893
scalar_td * td1
894-
895-
896-
class TestTimedeltaSeriesInvalidArithmeticOps(object):
897-
@pytest.mark.parametrize('scalar_td', [
898-
timedelta(minutes=5, seconds=4),
899-
Timedelta('5m4s'),
900-
Timedelta('5m4s').to_timedelta64()])
901-
def test_td64series_pow_invalid(self, scalar_td):
902-
td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
903-
td1.iloc[2] = np.nan
904-
905-
# check that we are getting a TypeError
906-
# with 'operate' (from core/ops.py) for the ops that are not
907-
# defined
908-
pattern = 'operate|unsupported|cannot|not supported'
909-
with tm.assert_raises_regex(TypeError, pattern):
910-
scalar_td ** td1
911-
with tm.assert_raises_regex(TypeError, pattern):
912-
td1 ** scalar_td

pandas/tests/test_arithmetic.py

+119
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
# -*- coding: utf-8 -*-
2+
# Arithmetc tests for DataFrame/Series/Index/Array classes that should
3+
# behave identically.
4+
from datetime import timedelta
5+
6+
import pytest
7+
import numpy as np
8+
9+
import pandas as pd
10+
import pandas.util.testing as tm
11+
12+
from pandas import Timedelta
13+
14+
15+
# ------------------------------------------------------------------
16+
# Numeric dtypes Arithmetic with Timedelta Scalar
17+
18+
class TestNumericArraylikeArithmeticWithTimedeltaScalar(object):
19+
20+
@pytest.mark.parametrize('box', [
21+
pd.Index,
22+
pd.Series,
23+
pytest.param(pd.DataFrame,
24+
marks=pytest.mark.xfail(reason="block.eval incorrect",
25+
strict=True))
26+
])
27+
@pytest.mark.parametrize('index', [
28+
pd.Int64Index(range(1, 11)),
29+
pd.UInt64Index(range(1, 11)),
30+
pd.Float64Index(range(1, 11)),
31+
pd.RangeIndex(1, 11)],
32+
ids=lambda x: type(x).__name__)
33+
@pytest.mark.parametrize('scalar_td', [
34+
Timedelta(days=1),
35+
Timedelta(days=1).to_timedelta64(),
36+
Timedelta(days=1).to_pytimedelta()],
37+
ids=lambda x: type(x).__name__)
38+
def test_index_mul_timedelta(self, scalar_td, index, box):
39+
# GH#19333
40+
41+
if (box is pd.Series and
42+
type(scalar_td) is timedelta and index.dtype == 'f8'):
43+
raise pytest.xfail(reason="Cannot multiply timedelta by float")
44+
45+
expected = pd.timedelta_range('1 days', '10 days')
46+
47+
index = tm.box_expected(index, box)
48+
expected = tm.box_expected(expected, box)
49+
50+
result = index * scalar_td
51+
tm.assert_equal(result, expected)
52+
53+
commute = scalar_td * index
54+
tm.assert_equal(commute, expected)
55+
56+
@pytest.mark.parametrize('box', [pd.Index, pd.Series, pd.DataFrame])
57+
@pytest.mark.parametrize('index', [
58+
pd.Int64Index(range(1, 3)),
59+
pd.UInt64Index(range(1, 3)),
60+
pd.Float64Index(range(1, 3)),
61+
pd.RangeIndex(1, 3)],
62+
ids=lambda x: type(x).__name__)
63+
@pytest.mark.parametrize('scalar_td', [
64+
Timedelta(days=1),
65+
Timedelta(days=1).to_timedelta64(),
66+
Timedelta(days=1).to_pytimedelta()],
67+
ids=lambda x: type(x).__name__)
68+
def test_index_rdiv_timedelta(self, scalar_td, index, box):
69+
70+
if box is pd.Series and type(scalar_td) is timedelta:
71+
raise pytest.xfail(reason="TODO: Figure out why this case fails")
72+
if box is pd.DataFrame and isinstance(scalar_td, timedelta):
73+
raise pytest.xfail(reason="TODO: Figure out why this case fails")
74+
75+
expected = pd.TimedeltaIndex(['1 Day', '12 Hours'])
76+
77+
index = tm.box_expected(index, box)
78+
expected = tm.box_expected(expected, box)
79+
80+
result = scalar_td / index
81+
tm.assert_equal(result, expected)
82+
83+
with pytest.raises(TypeError):
84+
index / scalar_td
85+
86+
87+
# ------------------------------------------------------------------
88+
# Timedelta64[ns] dtype Arithmetic Operations
89+
90+
91+
class TestTimedeltaArraylikeInvalidArithmeticOps(object):
92+
93+
@pytest.mark.parametrize('box', [
94+
pd.Index,
95+
pd.Series,
96+
pytest.param(pd.DataFrame,
97+
marks=pytest.mark.xfail(reason="raises ValueError "
98+
"instead of TypeError",
99+
strict=True))
100+
])
101+
@pytest.mark.parametrize('scalar_td', [
102+
timedelta(minutes=5, seconds=4),
103+
Timedelta('5m4s'),
104+
Timedelta('5m4s').to_timedelta64()])
105+
def test_td64series_pow_invalid(self, scalar_td, box):
106+
td1 = pd.Series([timedelta(minutes=5, seconds=3)] * 3)
107+
td1.iloc[2] = np.nan
108+
109+
td1 = tm.box_expected(td1, box)
110+
111+
# check that we are getting a TypeError
112+
# with 'operate' (from core/ops.py) for the ops that are not
113+
# defined
114+
pattern = 'operate|unsupported|cannot|not supported'
115+
with tm.assert_raises_regex(TypeError, pattern):
116+
scalar_td ** td1
117+
118+
with tm.assert_raises_regex(TypeError, pattern):
119+
td1 ** scalar_td

pandas/util/testing.py

+44
Original file line numberDiff line numberDiff line change
@@ -1478,6 +1478,50 @@ def assert_panel_equal(left, right,
14781478
assert item in left, msg
14791479

14801480

1481+
def assert_equal(left, right, **kwargs):
1482+
"""
1483+
Wrapper for tm.assert_*_equal to dispatch to the appropriate test function.
1484+
1485+
Parameters
1486+
----------
1487+
left : Index, Series, or DataFrame
1488+
right : Index, Series, or DataFrame
1489+
**kwargs
1490+
"""
1491+
if isinstance(left, pd.Index):
1492+
assert_index_equal(left, right, **kwargs)
1493+
elif isinstance(left, pd.Series):
1494+
assert_series_equal(left, right, **kwargs)
1495+
elif isinstance(left, pd.DataFrame):
1496+
assert_frame_equal(left, right, **kwargs)
1497+
else:
1498+
raise NotImplementedError(type(left))
1499+
1500+
1501+
def box_expected(expected, box_cls):
1502+
"""
1503+
Helper function to wrap the expected output of a test in a given box_class.
1504+
1505+
Parameters
1506+
----------
1507+
expected : np.ndarray, Index, Series
1508+
box_cls : {Index, Series, DataFrame}
1509+
1510+
Returns
1511+
-------
1512+
subclass of box_cls
1513+
"""
1514+
if box_cls is pd.Index:
1515+
expected = pd.Index(expected)
1516+
elif box_cls is pd.Series:
1517+
expected = pd.Series(expected)
1518+
elif box_cls is pd.DataFrame:
1519+
expected = pd.Series(expected).to_frame()
1520+
else:
1521+
raise NotImplementedError(box_cls)
1522+
return expected
1523+
1524+
14811525
# -----------------------------------------------------------------------------
14821526
# Sparse
14831527

0 commit comments

Comments
 (0)