Skip to content

Commit c9334fe

Browse files
jbrockmendelharisbal
authored and
harisbal
committed
Organize PeriodIndex tests (pandas-dev#19641)
1 parent 8bbb469 commit c9334fe

File tree

6 files changed

+358
-345
lines changed

6 files changed

+358
-345
lines changed

pandas/tests/indexes/period/test_arithmetic.py

+256
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,171 @@
1111
import pandas.core.indexes.period as period
1212

1313

14+
class TestPeriodIndexComparisons(object):
15+
@pytest.mark.parametrize('freq', ['M', '2M', '3M'])
16+
def test_pi_cmp_pi(self, freq):
17+
base = PeriodIndex(['2011-01', '2011-02', '2011-03', '2011-04'],
18+
freq=freq)
19+
per = Period('2011-02', freq=freq)
20+
21+
exp = np.array([False, True, False, False])
22+
tm.assert_numpy_array_equal(base == per, exp)
23+
tm.assert_numpy_array_equal(per == base, exp)
24+
25+
exp = np.array([True, False, True, True])
26+
tm.assert_numpy_array_equal(base != per, exp)
27+
tm.assert_numpy_array_equal(per != base, exp)
28+
29+
exp = np.array([False, False, True, True])
30+
tm.assert_numpy_array_equal(base > per, exp)
31+
tm.assert_numpy_array_equal(per < base, exp)
32+
33+
exp = np.array([True, False, False, False])
34+
tm.assert_numpy_array_equal(base < per, exp)
35+
tm.assert_numpy_array_equal(per > base, exp)
36+
37+
exp = np.array([False, True, True, True])
38+
tm.assert_numpy_array_equal(base >= per, exp)
39+
tm.assert_numpy_array_equal(per <= base, exp)
40+
41+
exp = np.array([True, True, False, False])
42+
tm.assert_numpy_array_equal(base <= per, exp)
43+
tm.assert_numpy_array_equal(per >= base, exp)
44+
45+
idx = PeriodIndex(['2011-02', '2011-01', '2011-03', '2011-05'],
46+
freq=freq)
47+
48+
exp = np.array([False, False, True, False])
49+
tm.assert_numpy_array_equal(base == idx, exp)
50+
51+
exp = np.array([True, True, False, True])
52+
tm.assert_numpy_array_equal(base != idx, exp)
53+
54+
exp = np.array([False, True, False, False])
55+
tm.assert_numpy_array_equal(base > idx, exp)
56+
57+
exp = np.array([True, False, False, True])
58+
tm.assert_numpy_array_equal(base < idx, exp)
59+
60+
exp = np.array([False, True, True, False])
61+
tm.assert_numpy_array_equal(base >= idx, exp)
62+
63+
exp = np.array([True, False, True, True])
64+
tm.assert_numpy_array_equal(base <= idx, exp)
65+
66+
@pytest.mark.parametrize('freq', ['M', '2M', '3M'])
67+
def test_pi_cmp_pi_mismatched_freq_raises(self, freq):
68+
# different base freq
69+
base = PeriodIndex(['2011-01', '2011-02', '2011-03', '2011-04'],
70+
freq=freq)
71+
72+
msg = "Input has different freq=A-DEC from PeriodIndex"
73+
with tm.assert_raises_regex(period.IncompatibleFrequency, msg):
74+
base <= Period('2011', freq='A')
75+
76+
with tm.assert_raises_regex(period.IncompatibleFrequency, msg):
77+
Period('2011', freq='A') >= base
78+
79+
idx = PeriodIndex(['2011', '2012', '2013', '2014'], freq='A')
80+
with tm.assert_raises_regex(period.IncompatibleFrequency, msg):
81+
base <= idx
82+
83+
# Different frequency
84+
msg = "Input has different freq=4M from PeriodIndex"
85+
with tm.assert_raises_regex(period.IncompatibleFrequency, msg):
86+
base <= Period('2011', freq='4M')
87+
88+
with tm.assert_raises_regex(period.IncompatibleFrequency, msg):
89+
Period('2011', freq='4M') >= base
90+
91+
idx = PeriodIndex(['2011', '2012', '2013', '2014'], freq='4M')
92+
with tm.assert_raises_regex(period.IncompatibleFrequency, msg):
93+
base <= idx
94+
95+
@pytest.mark.parametrize('freq', ['M', '2M', '3M'])
96+
def test_pi_cmp_nat(self, freq):
97+
idx1 = PeriodIndex(['2011-01', '2011-02', 'NaT', '2011-05'], freq=freq)
98+
99+
result = idx1 > Period('2011-02', freq=freq)
100+
exp = np.array([False, False, False, True])
101+
tm.assert_numpy_array_equal(result, exp)
102+
result = Period('2011-02', freq=freq) < idx1
103+
tm.assert_numpy_array_equal(result, exp)
104+
105+
result = idx1 == Period('NaT', freq=freq)
106+
exp = np.array([False, False, False, False])
107+
tm.assert_numpy_array_equal(result, exp)
108+
result = Period('NaT', freq=freq) == idx1
109+
tm.assert_numpy_array_equal(result, exp)
110+
111+
result = idx1 != Period('NaT', freq=freq)
112+
exp = np.array([True, True, True, True])
113+
tm.assert_numpy_array_equal(result, exp)
114+
result = Period('NaT', freq=freq) != idx1
115+
tm.assert_numpy_array_equal(result, exp)
116+
117+
idx2 = PeriodIndex(['2011-02', '2011-01', '2011-04', 'NaT'], freq=freq)
118+
result = idx1 < idx2
119+
exp = np.array([True, False, False, False])
120+
tm.assert_numpy_array_equal(result, exp)
121+
122+
result = idx1 == idx2
123+
exp = np.array([False, False, False, False])
124+
tm.assert_numpy_array_equal(result, exp)
125+
126+
result = idx1 != idx2
127+
exp = np.array([True, True, True, True])
128+
tm.assert_numpy_array_equal(result, exp)
129+
130+
result = idx1 == idx1
131+
exp = np.array([True, True, False, True])
132+
tm.assert_numpy_array_equal(result, exp)
133+
134+
result = idx1 != idx1
135+
exp = np.array([False, False, True, False])
136+
tm.assert_numpy_array_equal(result, exp)
137+
138+
@pytest.mark.parametrize('freq', ['M', '2M', '3M'])
139+
def test_pi_cmp_nat_mismatched_freq_raises(self, freq):
140+
idx1 = PeriodIndex(['2011-01', '2011-02', 'NaT', '2011-05'], freq=freq)
141+
142+
diff = PeriodIndex(['2011-02', '2011-01', '2011-04', 'NaT'], freq='4M')
143+
msg = "Input has different freq=4M from PeriodIndex"
144+
with tm.assert_raises_regex(period.IncompatibleFrequency, msg):
145+
idx1 > diff
146+
147+
with tm.assert_raises_regex(period.IncompatibleFrequency, msg):
148+
idx1 == diff
149+
150+
# TODO: De-duplicate with test_pi_cmp_nat
151+
def test_comp_nat(self):
152+
left = pd.PeriodIndex([pd.Period('2011-01-01'), pd.NaT,
153+
pd.Period('2011-01-03')])
154+
right = pd.PeriodIndex([pd.NaT, pd.NaT, pd.Period('2011-01-03')])
155+
156+
for lhs, rhs in [(left, right),
157+
(left.astype(object), right.astype(object))]:
158+
result = lhs == rhs
159+
expected = np.array([False, False, True])
160+
tm.assert_numpy_array_equal(result, expected)
161+
162+
result = lhs != rhs
163+
expected = np.array([True, True, False])
164+
tm.assert_numpy_array_equal(result, expected)
165+
166+
expected = np.array([False, False, False])
167+
tm.assert_numpy_array_equal(lhs == pd.NaT, expected)
168+
tm.assert_numpy_array_equal(pd.NaT == rhs, expected)
169+
170+
expected = np.array([True, True, True])
171+
tm.assert_numpy_array_equal(lhs != pd.NaT, expected)
172+
tm.assert_numpy_array_equal(pd.NaT != lhs, expected)
173+
174+
expected = np.array([False, False, False])
175+
tm.assert_numpy_array_equal(lhs < pd.NaT, expected)
176+
tm.assert_numpy_array_equal(pd.NaT > lhs, expected)
177+
178+
14179
class TestPeriodIndexArithmetic(object):
15180
def test_pi_add_offset_array(self):
16181
# GH#18849
@@ -250,6 +415,97 @@ def test_sub_isub(self):
250415
rng -= 1
251416
tm.assert_index_equal(rng, expected)
252417

418+
# ---------------------------------------------------------------
419+
# PeriodIndex.shift is used by __add__ and __sub__
420+
421+
def test_pi_shift_ndarray(self):
422+
idx = PeriodIndex(['2011-01', '2011-02', 'NaT',
423+
'2011-04'], freq='M', name='idx')
424+
result = idx.shift(np.array([1, 2, 3, 4]))
425+
expected = PeriodIndex(['2011-02', '2011-04', 'NaT',
426+
'2011-08'], freq='M', name='idx')
427+
tm.assert_index_equal(result, expected)
428+
429+
idx = PeriodIndex(['2011-01', '2011-02', 'NaT',
430+
'2011-04'], freq='M', name='idx')
431+
result = idx.shift(np.array([1, -2, 3, -4]))
432+
expected = PeriodIndex(['2011-02', '2010-12', 'NaT',
433+
'2010-12'], freq='M', name='idx')
434+
tm.assert_index_equal(result, expected)
435+
436+
def test_shift(self):
437+
pi1 = PeriodIndex(freq='A', start='1/1/2001', end='12/1/2009')
438+
pi2 = PeriodIndex(freq='A', start='1/1/2002', end='12/1/2010')
439+
440+
tm.assert_index_equal(pi1.shift(0), pi1)
441+
442+
assert len(pi1) == len(pi2)
443+
tm.assert_index_equal(pi1.shift(1), pi2)
444+
445+
pi1 = PeriodIndex(freq='A', start='1/1/2001', end='12/1/2009')
446+
pi2 = PeriodIndex(freq='A', start='1/1/2000', end='12/1/2008')
447+
assert len(pi1) == len(pi2)
448+
tm.assert_index_equal(pi1.shift(-1), pi2)
449+
450+
pi1 = PeriodIndex(freq='M', start='1/1/2001', end='12/1/2009')
451+
pi2 = PeriodIndex(freq='M', start='2/1/2001', end='1/1/2010')
452+
assert len(pi1) == len(pi2)
453+
tm.assert_index_equal(pi1.shift(1), pi2)
454+
455+
pi1 = PeriodIndex(freq='M', start='1/1/2001', end='12/1/2009')
456+
pi2 = PeriodIndex(freq='M', start='12/1/2000', end='11/1/2009')
457+
assert len(pi1) == len(pi2)
458+
tm.assert_index_equal(pi1.shift(-1), pi2)
459+
460+
pi1 = PeriodIndex(freq='D', start='1/1/2001', end='12/1/2009')
461+
pi2 = PeriodIndex(freq='D', start='1/2/2001', end='12/2/2009')
462+
assert len(pi1) == len(pi2)
463+
tm.assert_index_equal(pi1.shift(1), pi2)
464+
465+
pi1 = PeriodIndex(freq='D', start='1/1/2001', end='12/1/2009')
466+
pi2 = PeriodIndex(freq='D', start='12/31/2000', end='11/30/2009')
467+
assert len(pi1) == len(pi2)
468+
tm.assert_index_equal(pi1.shift(-1), pi2)
469+
470+
def test_shift_corner_cases(self):
471+
# GH#9903
472+
idx = pd.PeriodIndex([], name='xxx', freq='H')
473+
474+
with pytest.raises(TypeError):
475+
# period shift doesn't accept freq
476+
idx.shift(1, freq='H')
477+
478+
tm.assert_index_equal(idx.shift(0), idx)
479+
tm.assert_index_equal(idx.shift(3), idx)
480+
481+
idx = pd.PeriodIndex(['2011-01-01 10:00', '2011-01-01 11:00'
482+
'2011-01-01 12:00'], name='xxx', freq='H')
483+
tm.assert_index_equal(idx.shift(0), idx)
484+
exp = pd.PeriodIndex(['2011-01-01 13:00', '2011-01-01 14:00'
485+
'2011-01-01 15:00'], name='xxx', freq='H')
486+
tm.assert_index_equal(idx.shift(3), exp)
487+
exp = pd.PeriodIndex(['2011-01-01 07:00', '2011-01-01 08:00'
488+
'2011-01-01 09:00'], name='xxx', freq='H')
489+
tm.assert_index_equal(idx.shift(-3), exp)
490+
491+
def test_shift_nat(self):
492+
idx = PeriodIndex(['2011-01', '2011-02', 'NaT',
493+
'2011-04'], freq='M', name='idx')
494+
result = idx.shift(1)
495+
expected = PeriodIndex(['2011-02', '2011-03', 'NaT',
496+
'2011-05'], freq='M', name='idx')
497+
tm.assert_index_equal(result, expected)
498+
assert result.name == expected.name
499+
500+
def test_shift_gh8083(self):
501+
# test shift for PeriodIndex
502+
# GH#8083
503+
drange = pd.period_range('20130101', periods=5, freq='D')
504+
result = drange.shift(1)
505+
expected = PeriodIndex(['2013-01-02', '2013-01-03', '2013-01-04',
506+
'2013-01-05', '2013-01-06'], freq='D')
507+
tm.assert_index_equal(result, expected)
508+
253509

254510
class TestPeriodIndexSeriesMethods(object):
255511
""" Test PeriodIndex and Period Series Ops consistency """

0 commit comments

Comments
 (0)