Skip to content

Commit 098cd52

Browse files
jbrockmendeljbrockmendel
authored
REF/TST: collect indexing tests by method (#37677)
1 parent fc38f46 commit 098cd52

17 files changed

+374
-347
lines changed

‎pandas/tests/frame/indexing/test_categorical.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -352,7 +352,7 @@ def test_assigning_ops(self):
352352
df.loc[2:3, "b"] = Categorical(["b", "b"], categories=["a", "b"])
353353
tm.assert_frame_equal(df, exp)
354354

355-
def test_setitem_single_row_categorical(self):
355+
def test_loc_setitem_single_row_categorical(self):
356356
# GH 25495
357357
df = DataFrame({"Alpha": ["a"], "Numeric": [0]})
358358
categories = Categorical(df["Alpha"], categories=["a", "b", "c"])

‎pandas/tests/frame/indexing/test_getitem.py

+106-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,16 @@
11
import numpy as np
22
import pytest
33

4-
from pandas import DataFrame, MultiIndex, period_range
4+
from pandas import (
5+
Categorical,
6+
CategoricalDtype,
7+
CategoricalIndex,
8+
DataFrame,
9+
MultiIndex,
10+
Timestamp,
11+
get_dummies,
12+
period_range,
13+
)
514
import pandas._testing as tm
615

716

@@ -29,3 +38,99 @@ def test_getitem_periodindex(self):
2938

3039
ts = df["1/1/2000"]
3140
tm.assert_series_equal(ts, df.iloc[:, 0])
41+
42+
def test_getitem_list_of_labels_categoricalindex_cols(self):
43+
# GH#16115
44+
cats = Categorical([Timestamp("12-31-1999"), Timestamp("12-31-2000")])
45+
46+
expected = DataFrame(
47+
[[1, 0], [0, 1]], dtype="uint8", index=[0, 1], columns=cats
48+
)
49+
dummies = get_dummies(cats)
50+
result = dummies[list(dummies.columns)]
51+
tm.assert_frame_equal(result, expected)
52+
53+
54+
class TestGetitemCallable:
55+
def test_getitem_callable(self, float_frame):
56+
# GH#12533
57+
result = float_frame[lambda x: "A"]
58+
expected = float_frame.loc[:, "A"]
59+
tm.assert_series_equal(result, expected)
60+
61+
result = float_frame[lambda x: ["A", "B"]]
62+
expected = float_frame.loc[:, ["A", "B"]]
63+
tm.assert_frame_equal(result, float_frame.loc[:, ["A", "B"]])
64+
65+
df = float_frame[:3]
66+
result = df[lambda x: [True, False, True]]
67+
expected = float_frame.iloc[[0, 2], :]
68+
tm.assert_frame_equal(result, expected)
69+
70+
71+
class TestGetitemBooleanMask:
72+
def test_getitem_bool_mask_categorical_index(self):
73+
74+
df3 = DataFrame(
75+
{
76+
"A": np.arange(6, dtype="int64"),
77+
},
78+
index=CategoricalIndex(
79+
[1, 1, 2, 1, 3, 2],
80+
dtype=CategoricalDtype([3, 2, 1], ordered=True),
81+
name="B",
82+
),
83+
)
84+
df4 = DataFrame(
85+
{
86+
"A": np.arange(6, dtype="int64"),
87+
},
88+
index=CategoricalIndex(
89+
[1, 1, 2, 1, 3, 2],
90+
dtype=CategoricalDtype([3, 2, 1], ordered=False),
91+
name="B",
92+
),
93+
)
94+
95+
result = df3[df3.index == "a"]
96+
expected = df3.iloc[[]]
97+
tm.assert_frame_equal(result, expected)
98+
99+
result = df4[df4.index == "a"]
100+
expected = df4.iloc[[]]
101+
tm.assert_frame_equal(result, expected)
102+
103+
result = df3[df3.index == 1]
104+
expected = df3.iloc[[0, 1, 3]]
105+
tm.assert_frame_equal(result, expected)
106+
107+
result = df4[df4.index == 1]
108+
expected = df4.iloc[[0, 1, 3]]
109+
tm.assert_frame_equal(result, expected)
110+
111+
# since we have an ordered categorical
112+
113+
# CategoricalIndex([1, 1, 2, 1, 3, 2],
114+
# categories=[3, 2, 1],
115+
# ordered=True,
116+
# name='B')
117+
result = df3[df3.index < 2]
118+
expected = df3.iloc[[4]]
119+
tm.assert_frame_equal(result, expected)
120+
121+
result = df3[df3.index > 1]
122+
expected = df3.iloc[[]]
123+
tm.assert_frame_equal(result, expected)
124+
125+
# unordered
126+
# cannot be compared
127+
128+
# CategoricalIndex([1, 1, 2, 1, 3, 2],
129+
# categories=[3, 2, 1],
130+
# ordered=False,
131+
# name='B')
132+
msg = "Unordered Categoricals can only compare equality or not"
133+
with pytest.raises(TypeError, match=msg):
134+
df4[df4.index < 2]
135+
with pytest.raises(TypeError, match=msg):
136+
df4[df4.index > 1]

‎pandas/tests/frame/indexing/test_indexing.py

-15
Original file line numberDiff line numberDiff line change
@@ -110,21 +110,6 @@ def test_getitem_listlike(self, idx_type, levels, float_frame):
110110
with pytest.raises(KeyError, match="not in index"):
111111
frame[idx]
112112

113-
def test_getitem_callable(self, float_frame):
114-
# GH 12533
115-
result = float_frame[lambda x: "A"]
116-
expected = float_frame.loc[:, "A"]
117-
tm.assert_series_equal(result, expected)
118-
119-
result = float_frame[lambda x: ["A", "B"]]
120-
expected = float_frame.loc[:, ["A", "B"]]
121-
tm.assert_frame_equal(result, float_frame.loc[:, ["A", "B"]])
122-
123-
df = float_frame[:3]
124-
result = df[lambda x: [True, False, True]]
125-
expected = float_frame.iloc[[0, 2], :]
126-
tm.assert_frame_equal(result, expected)
127-
128113
def test_setitem_list(self, float_frame):
129114

130115
float_frame["E"] = "foo"

‎pandas/tests/indexing/test_at.py

+70-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import numpy as np
44
import pytest
55

6-
from pandas import DataFrame
6+
from pandas import DataFrame, Series
77
import pandas._testing as tm
88

99

@@ -39,3 +39,72 @@ def test_at_with_duplicate_axes_requires_scalar_lookup(self):
3939
df.at[1, ["A"]] = 1
4040
with pytest.raises(ValueError, match=msg):
4141
df.at[:, "A"] = 1
42+
43+
44+
class TestAtErrors:
45+
# TODO: De-duplicate/parametrize
46+
# test_at_series_raises_key_error, test_at_frame_raises_key_error,
47+
# test_at_series_raises_key_error2, test_at_frame_raises_key_error2
48+
49+
def test_at_series_raises_key_error(self):
50+
# GH#31724 .at should match .loc
51+
52+
ser = Series([1, 2, 3], index=[3, 2, 1])
53+
result = ser.at[1]
54+
assert result == 3
55+
result = ser.loc[1]
56+
assert result == 3
57+
58+
with pytest.raises(KeyError, match="a"):
59+
ser.at["a"]
60+
with pytest.raises(KeyError, match="a"):
61+
# .at should match .loc
62+
ser.loc["a"]
63+
64+
def test_at_frame_raises_key_error(self):
65+
# GH#31724 .at should match .loc
66+
67+
df = DataFrame({0: [1, 2, 3]}, index=[3, 2, 1])
68+
69+
result = df.at[1, 0]
70+
assert result == 3
71+
result = df.loc[1, 0]
72+
assert result == 3
73+
74+
with pytest.raises(KeyError, match="a"):
75+
df.at["a", 0]
76+
with pytest.raises(KeyError, match="a"):
77+
df.loc["a", 0]
78+
79+
with pytest.raises(KeyError, match="a"):
80+
df.at[1, "a"]
81+
with pytest.raises(KeyError, match="a"):
82+
df.loc[1, "a"]
83+
84+
def test_at_series_raises_key_error2(self):
85+
# at should not fallback
86+
# GH#7814
87+
# GH#31724 .at should match .loc
88+
ser = Series([1, 2, 3], index=list("abc"))
89+
result = ser.at["a"]
90+
assert result == 1
91+
result = ser.loc["a"]
92+
assert result == 1
93+
94+
with pytest.raises(KeyError, match="^0$"):
95+
ser.at[0]
96+
with pytest.raises(KeyError, match="^0$"):
97+
ser.loc[0]
98+
99+
def test_at_frame_raises_key_error2(self):
100+
# GH#31724 .at should match .loc
101+
df = DataFrame({"A": [1, 2, 3]}, index=list("abc"))
102+
result = df.at["a", "A"]
103+
assert result == 1
104+
result = df.loc["a", "A"]
105+
assert result == 1
106+
107+
with pytest.raises(KeyError, match="^0$"):
108+
df.at["a", 0]
109+
with pytest.raises(KeyError, match="^0$"):
110+
df.loc["a", 0]

‎pandas/tests/indexing/test_categorical.py

-95
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
import pytest
33

44
from pandas.core.dtypes.common import is_categorical_dtype
5-
from pandas.core.dtypes.dtypes import CategoricalDtype
65

76
import pandas as pd
87
from pandas import (
@@ -276,27 +275,6 @@ def test_slicing_doc_examples(self):
276275
)
277276
tm.assert_frame_equal(result, expected)
278277

279-
def test_getitem_category_type(self):
280-
# GH 14580
281-
# test iloc() on Series with Categorical data
282-
283-
s = Series([1, 2, 3]).astype("category")
284-
285-
# get slice
286-
result = s.iloc[0:2]
287-
expected = Series([1, 2]).astype(CategoricalDtype([1, 2, 3]))
288-
tm.assert_series_equal(result, expected)
289-
290-
# get list of indexes
291-
result = s.iloc[[0, 1]]
292-
expected = Series([1, 2]).astype(CategoricalDtype([1, 2, 3]))
293-
tm.assert_series_equal(result, expected)
294-
295-
# get boolean array
296-
result = s.iloc[[True, False, False]]
297-
expected = Series([1]).astype(CategoricalDtype([1, 2, 3]))
298-
tm.assert_series_equal(result, expected)
299-
300278
def test_loc_listlike(self):
301279

302280
# list of labels
@@ -413,17 +391,6 @@ def test_loc_listlike_dtypes(self):
413391
with pytest.raises(KeyError, match=msg):
414392
df.loc[["a", "x"]]
415393

416-
def test_getitem_with_listlike(self):
417-
# GH 16115
418-
cats = Categorical([Timestamp("12-31-1999"), Timestamp("12-31-2000")])
419-
420-
expected = DataFrame(
421-
[[1, 0], [0, 1]], dtype="uint8", index=[0, 1], columns=cats
422-
)
423-
dummies = pd.get_dummies(cats)
424-
result = dummies[list(dummies.columns)]
425-
tm.assert_frame_equal(result, expected)
426-
427394
def test_ix_categorical_index(self):
428395
# GH 12531
429396
df = DataFrame(np.random.randn(3, 3), index=list("ABC"), columns=list("XYZ"))
@@ -512,68 +479,6 @@ def test_loc_and_at_with_categorical_index(self):
512479
assert df.loc["B", 1] == 4
513480
assert df.at["B", 1] == 4
514481

515-
def test_getitem_bool_mask_categorical_index(self):
516-
517-
df3 = DataFrame(
518-
{
519-
"A": np.arange(6, dtype="int64"),
520-
},
521-
index=CategoricalIndex(
522-
[1, 1, 2, 1, 3, 2], dtype=CDT([3, 2, 1], ordered=True), name="B"
523-
),
524-
)
525-
df4 = DataFrame(
526-
{
527-
"A": np.arange(6, dtype="int64"),
528-
},
529-
index=CategoricalIndex(
530-
[1, 1, 2, 1, 3, 2], dtype=CDT([3, 2, 1], ordered=False), name="B"
531-
),
532-
)
533-
534-
result = df3[df3.index == "a"]
535-
expected = df3.iloc[[]]
536-
tm.assert_frame_equal(result, expected)
537-
538-
result = df4[df4.index == "a"]
539-
expected = df4.iloc[[]]
540-
tm.assert_frame_equal(result, expected)
541-
542-
result = df3[df3.index == 1]
543-
expected = df3.iloc[[0, 1, 3]]
544-
tm.assert_frame_equal(result, expected)
545-
546-
result = df4[df4.index == 1]
547-
expected = df4.iloc[[0, 1, 3]]
548-
tm.assert_frame_equal(result, expected)
549-
550-
# since we have an ordered categorical
551-
552-
# CategoricalIndex([1, 1, 2, 1, 3, 2],
553-
# categories=[3, 2, 1],
554-
# ordered=True,
555-
# name='B')
556-
result = df3[df3.index < 2]
557-
expected = df3.iloc[[4]]
558-
tm.assert_frame_equal(result, expected)
559-
560-
result = df3[df3.index > 1]
561-
expected = df3.iloc[[]]
562-
tm.assert_frame_equal(result, expected)
563-
564-
# unordered
565-
# cannot be compared
566-
567-
# CategoricalIndex([1, 1, 2, 1, 3, 2],
568-
# categories=[3, 2, 1],
569-
# ordered=False,
570-
# name='B')
571-
msg = "Unordered Categoricals can only compare equality or not"
572-
with pytest.raises(TypeError, match=msg):
573-
df4[df4.index < 2]
574-
with pytest.raises(TypeError, match=msg):
575-
df4[df4.index > 1]
576-
577482
def test_indexing_with_category(self):
578483

579484
# https://github.com/pandas-dev/pandas/issues/12564

0 commit comments

Comments
 (0)