|
7 | 7 |
|
8 | 8 |
|
9 | 9 | class TestIntervalIndex:
|
10 |
| - def setup_method(self, method): |
11 |
| - self.s = Series(np.arange(5), IntervalIndex.from_breaks(np.arange(6))) |
| 10 | + @pytest.fixture |
| 11 | + def series_with_interval_index(self): |
| 12 | + return Series(np.arange(5), IntervalIndex.from_breaks(np.arange(6))) |
12 | 13 |
|
13 |
| - def test_getitem_with_scalar(self): |
| 14 | + def test_getitem_with_scalar(self, series_with_interval_index, indexer_sl): |
14 | 15 |
|
15 |
| - s = self.s |
| 16 | + ser = series_with_interval_index.copy() |
16 | 17 |
|
17 |
| - expected = s.iloc[:3] |
18 |
| - tm.assert_series_equal(expected, s[:3]) |
19 |
| - tm.assert_series_equal(expected, s[:2.5]) |
20 |
| - tm.assert_series_equal(expected, s[0.1:2.5]) |
| 18 | + expected = ser.iloc[:3] |
| 19 | + tm.assert_series_equal(expected, indexer_sl(ser)[:3]) |
| 20 | + tm.assert_series_equal(expected, indexer_sl(ser)[:2.5]) |
| 21 | + tm.assert_series_equal(expected, indexer_sl(ser)[0.1:2.5]) |
| 22 | + if indexer_sl is tm.loc: |
| 23 | + tm.assert_series_equal(expected, ser.loc[-1:3]) |
21 | 24 |
|
22 |
| - expected = s.iloc[1:4] |
23 |
| - tm.assert_series_equal(expected, s[[1.5, 2.5, 3.5]]) |
24 |
| - tm.assert_series_equal(expected, s[[2, 3, 4]]) |
25 |
| - tm.assert_series_equal(expected, s[[1.5, 3, 4]]) |
| 25 | + expected = ser.iloc[1:4] |
| 26 | + tm.assert_series_equal(expected, indexer_sl(ser)[[1.5, 2.5, 3.5]]) |
| 27 | + tm.assert_series_equal(expected, indexer_sl(ser)[[2, 3, 4]]) |
| 28 | + tm.assert_series_equal(expected, indexer_sl(ser)[[1.5, 3, 4]]) |
26 | 29 |
|
27 |
| - expected = s.iloc[2:5] |
28 |
| - tm.assert_series_equal(expected, s[s >= 2]) |
| 30 | + expected = ser.iloc[2:5] |
| 31 | + tm.assert_series_equal(expected, indexer_sl(ser)[ser >= 2]) |
29 | 32 |
|
30 | 33 | @pytest.mark.parametrize("direction", ["increasing", "decreasing"])
|
31 |
| - def test_nonoverlapping_monotonic(self, direction, closed): |
| 34 | + def test_nonoverlapping_monotonic(self, direction, closed, indexer_sl): |
32 | 35 | tpls = [(0, 1), (2, 3), (4, 5)]
|
33 | 36 | if direction == "decreasing":
|
34 | 37 | tpls = tpls[::-1]
|
35 | 38 |
|
36 | 39 | idx = IntervalIndex.from_tuples(tpls, closed=closed)
|
37 |
| - s = Series(list("abc"), idx) |
| 40 | + ser = Series(list("abc"), idx) |
38 | 41 |
|
39 |
| - for key, expected in zip(idx.left, s): |
| 42 | + for key, expected in zip(idx.left, ser): |
40 | 43 | if idx.closed_left:
|
41 |
| - assert s[key] == expected |
42 |
| - assert s.loc[key] == expected |
| 44 | + assert indexer_sl(ser)[key] == expected |
43 | 45 | else:
|
44 | 46 | with pytest.raises(KeyError, match=str(key)):
|
45 |
| - s[key] |
46 |
| - with pytest.raises(KeyError, match=str(key)): |
47 |
| - s.loc[key] |
| 47 | + indexer_sl(ser)[key] |
48 | 48 |
|
49 |
| - for key, expected in zip(idx.right, s): |
| 49 | + for key, expected in zip(idx.right, ser): |
50 | 50 | if idx.closed_right:
|
51 |
| - assert s[key] == expected |
52 |
| - assert s.loc[key] == expected |
| 51 | + assert indexer_sl(ser)[key] == expected |
53 | 52 | else:
|
54 | 53 | with pytest.raises(KeyError, match=str(key)):
|
55 |
| - s[key] |
56 |
| - with pytest.raises(KeyError, match=str(key)): |
57 |
| - s.loc[key] |
| 54 | + indexer_sl(ser)[key] |
58 | 55 |
|
59 |
| - for key, expected in zip(idx.mid, s): |
60 |
| - assert s[key] == expected |
61 |
| - assert s.loc[key] == expected |
| 56 | + for key, expected in zip(idx.mid, ser): |
| 57 | + assert indexer_sl(ser)[key] == expected |
62 | 58 |
|
63 |
| - def test_non_matching(self): |
64 |
| - s = self.s |
| 59 | + def test_non_matching(self, series_with_interval_index, indexer_sl): |
| 60 | + ser = series_with_interval_index.copy() |
65 | 61 |
|
66 | 62 | # this is a departure from our current
|
67 | 63 | # indexing scheme, but simpler
|
68 | 64 | with pytest.raises(KeyError, match=r"^\[-1\]$"):
|
69 |
| - s.loc[[-1, 3, 4, 5]] |
| 65 | + indexer_sl(ser)[[-1, 3, 4, 5]] |
70 | 66 |
|
71 | 67 | with pytest.raises(KeyError, match=r"^\[-1\]$"):
|
72 |
| - s.loc[[-1, 3]] |
| 68 | + indexer_sl(ser)[[-1, 3]] |
73 | 69 |
|
74 | 70 | @pytest.mark.arm_slow
|
75 | 71 | def test_large_series(self):
|
76 |
| - s = Series( |
| 72 | + ser = Series( |
77 | 73 | np.arange(1000000), index=IntervalIndex.from_breaks(np.arange(1000001))
|
78 | 74 | )
|
79 | 75 |
|
80 |
| - result1 = s.loc[:80000] |
81 |
| - result2 = s.loc[0:80000] |
82 |
| - result3 = s.loc[0:80000:1] |
| 76 | + result1 = ser.loc[:80000] |
| 77 | + result2 = ser.loc[0:80000] |
| 78 | + result3 = ser.loc[0:80000:1] |
83 | 79 | tm.assert_series_equal(result1, result2)
|
84 | 80 | tm.assert_series_equal(result1, result3)
|
85 | 81 |
|
86 | 82 | def test_loc_getitem_frame(self):
|
87 | 83 | # CategoricalIndex with IntervalIndex categories
|
88 | 84 | df = DataFrame({"A": range(10)})
|
89 |
| - s = pd.cut(df.A, 5) |
90 |
| - df["B"] = s |
| 85 | + ser = pd.cut(df.A, 5) |
| 86 | + df["B"] = ser |
91 | 87 | df = df.set_index("B")
|
92 | 88 |
|
93 | 89 | result = df.loc[4]
|
|
0 commit comments