|
11 | 11 | import pandas.core.indexes.period as period
|
12 | 12 |
|
13 | 13 |
|
| 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 | + |
14 | 179 | class TestPeriodIndexArithmetic(object):
|
15 | 180 | def test_pi_add_offset_array(self):
|
16 | 181 | # GH#18849
|
@@ -250,6 +415,97 @@ def test_sub_isub(self):
|
250 | 415 | rng -= 1
|
251 | 416 | tm.assert_index_equal(rng, expected)
|
252 | 417 |
|
| 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 | + |
253 | 509 |
|
254 | 510 | class TestPeriodIndexSeriesMethods(object):
|
255 | 511 | """ Test PeriodIndex and Period Series Ops consistency """
|
|
0 commit comments