|
7 | 7 | @pytest.mark.parametrize(
|
8 | 8 | "values, dtype",
|
9 | 9 | [
|
| 10 | + ([], "object"), |
10 | 11 | ([1, 2, 3], "int64"),
|
11 | 12 | ([1.0, 2.0, 3.0], "float64"),
|
12 | 13 | (["a", "b", "c"], "object"),
|
|
22 | 23 | @pytest.mark.parametrize(
|
23 | 24 | "mask", [[True, False, False], [True, True, True], [False, False, False]]
|
24 | 25 | )
|
25 |
| -@pytest.mark.parametrize("box_mask", [True, False]) |
| 26 | +@pytest.mark.parametrize("indexer_class", [list, pd.array, pd.Index, pd.Series]) |
26 | 27 | @pytest.mark.parametrize("frame", [True, False])
|
27 |
| -def test_series_mask_boolean(values, dtype, mask, box_mask, frame): |
28 |
| - ser = pd.Series(values, dtype=dtype, index=["a", "b", "c"]) |
29 |
| - if frame: |
30 |
| - ser = ser.to_frame() |
31 |
| - mask = pd.array(mask, dtype="boolean") |
32 |
| - if box_mask: |
33 |
| - mask = pd.Series(mask, index=ser.index) |
34 |
| - |
35 |
| - expected = ser[mask.astype("bool")] |
| 28 | +def test_series_mask_boolean(values, dtype, mask, indexer_class, frame): |
| 29 | + # In case len(values) < 3 |
| 30 | + index = ["a", "b", "c"][: len(values)] |
| 31 | + mask = mask[: len(values)] |
36 | 32 |
|
37 |
| - result = ser[mask] |
38 |
| - tm.assert_equal(result, expected) |
39 |
| - |
40 |
| - if not box_mask: |
41 |
| - # Series.iloc[Series[bool]] isn't allowed |
42 |
| - result = ser.iloc[mask] |
43 |
| - tm.assert_equal(result, expected) |
| 33 | + obj = pd.Series(values, dtype=dtype, index=index) |
| 34 | + if frame: |
| 35 | + if len(values) == 0: |
| 36 | + # Otherwise obj is an empty DataFrame with shape (0, 1) |
| 37 | + obj = pd.DataFrame(dtype=dtype) |
| 38 | + else: |
| 39 | + obj = obj.to_frame() |
| 40 | + |
| 41 | + if indexer_class is pd.array: |
| 42 | + mask = pd.array(mask, dtype="boolean") |
| 43 | + elif indexer_class is pd.Series: |
| 44 | + mask = pd.Series(mask, index=obj.index, dtype="boolean") |
| 45 | + else: |
| 46 | + mask = indexer_class(mask) |
44 | 47 |
|
45 |
| - result = ser.loc[mask] |
46 |
| - tm.assert_equal(result, expected) |
| 48 | + expected = obj[mask] |
47 | 49 |
|
48 |
| - # empty |
49 |
| - mask = mask[:0] |
50 |
| - ser = ser.iloc[:0] |
51 |
| - expected = ser[mask.astype("bool")] |
52 |
| - result = ser[mask] |
| 50 | + result = obj[mask] |
53 | 51 | tm.assert_equal(result, expected)
|
54 | 52 |
|
55 |
| - if not box_mask: |
56 |
| - # Series.iloc[Series[bool]] isn't allowed |
57 |
| - result = ser.iloc[mask] |
| 53 | + if indexer_class is pd.Series: |
| 54 | + msg = "iLocation based boolean indexing cannot use an indexable as a mask" |
| 55 | + with pytest.raises(ValueError, match=msg): |
| 56 | + result = obj.iloc[mask] |
| 57 | + tm.assert_equal(result, expected) |
| 58 | + else: |
| 59 | + result = obj.iloc[mask] |
58 | 60 | tm.assert_equal(result, expected)
|
59 | 61 |
|
60 |
| - result = ser.loc[mask] |
| 62 | + result = obj.loc[mask] |
61 | 63 | tm.assert_equal(result, expected)
|
62 | 64 |
|
63 | 65 |
|
|
0 commit comments