1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
| import pytest
|
| import pandas as pd
| import pandas._testing as tm
|
|
| @pytest.mark.parametrize(
| "values, dtype",
| [
| ([], "object"),
| ([1, 2, 3], "int64"),
| ([1.0, 2.0, 3.0], "float64"),
| (["a", "b", "c"], "object"),
| (["a", "b", "c"], "string"),
| ([1, 2, 3], "datetime64[ns]"),
| ([1, 2, 3], "datetime64[ns, CET]"),
| ([1, 2, 3], "timedelta64[ns]"),
| (["2000", "2001", "2002"], "Period[D]"),
| ([1, 0, 3], "Sparse"),
| ([pd.Interval(0, 1), pd.Interval(1, 2), pd.Interval(3, 4)], "interval"),
| ],
| )
| @pytest.mark.parametrize(
| "mask", [[True, False, False], [True, True, True], [False, False, False]]
| )
| @pytest.mark.parametrize("indexer_class", [list, pd.array, pd.Index, pd.Series])
| @pytest.mark.parametrize("frame", [True, False])
| def test_series_mask_boolean(values, dtype, mask, indexer_class, frame):
| # In case len(values) < 3
| index = ["a", "b", "c"][: len(values)]
| mask = mask[: len(values)]
|
| obj = pd.Series(values, dtype=dtype, index=index)
| if frame:
| if len(values) == 0:
| # Otherwise obj is an empty DataFrame with shape (0, 1)
| obj = pd.DataFrame(dtype=dtype, index=index)
| else:
| obj = obj.to_frame()
|
| if indexer_class is pd.array:
| mask = pd.array(mask, dtype="boolean")
| elif indexer_class is pd.Series:
| mask = pd.Series(mask, index=obj.index, dtype="boolean")
| else:
| mask = indexer_class(mask)
|
| expected = obj[mask]
|
| result = obj[mask]
| tm.assert_equal(result, expected)
|
| if indexer_class is pd.Series:
| msg = "iLocation based boolean indexing cannot use an indexable as a mask"
| with pytest.raises(ValueError, match=msg):
| result = obj.iloc[mask]
| tm.assert_equal(result, expected)
| else:
| result = obj.iloc[mask]
| tm.assert_equal(result, expected)
|
| result = obj.loc[mask]
| tm.assert_equal(result, expected)
|
|
| def test_na_treated_as_false(frame_or_series, indexer_sli):
| # https://github.com/pandas-dev/pandas/issues/31503
| obj = frame_or_series([1, 2, 3])
|
| mask = pd.array([True, False, None], dtype="boolean")
|
| result = indexer_sli(obj)[mask]
| expected = indexer_sli(obj)[mask.fillna(False)]
|
| tm.assert_equal(result, expected)
|
|
