REF/TST: collect indexing tests by method (#37677)

Author: jbrockmendel

pandas/tests/frame/indexing/test_categorical.py

@@ -352,7 +352,7 @@ def test_assigning_ops(self):
         df.loc[2:3, "b"] = Categorical(["b", "b"], categories=["a", "b"])
         tm.assert_frame_equal(df, exp)
 
-    def test_setitem_single_row_categorical(self):
+    def test_loc_setitem_single_row_categorical(self):
         # GH 25495
         df = DataFrame({"Alpha": ["a"], "Numeric": [0]})
         categories = Categorical(df["Alpha"], categories=["a", "b", "c"])

pandas/tests/frame/indexing/test_getitem.py

@@ -1,7 +1,16 @@
 import numpy as np
 import pytest
 
-from pandas import DataFrame, MultiIndex, period_range
+from pandas import (
+    Categorical,
+    CategoricalDtype,
+    CategoricalIndex,
+    DataFrame,
+    MultiIndex,
+    Timestamp,
+    get_dummies,
+    period_range,
+)
 import pandas._testing as tm
 
 
@@ -29,3 +38,99 @@ def test_getitem_periodindex(self):
 
         ts = df["1/1/2000"]
         tm.assert_series_equal(ts, df.iloc[:, 0])
+
+    def test_getitem_list_of_labels_categoricalindex_cols(self):
+        # GH#16115
+        cats = Categorical([Timestamp("12-31-1999"), Timestamp("12-31-2000")])
+
+        expected = DataFrame(
+            [[1, 0], [0, 1]], dtype="uint8", index=[0, 1], columns=cats
+        )
+        dummies = get_dummies(cats)
+        result = dummies[list(dummies.columns)]
+        tm.assert_frame_equal(result, expected)
+
+
+class TestGetitemCallable:
+    def test_getitem_callable(self, float_frame):
+        # GH#12533
+        result = float_frame[lambda x: "A"]
+        expected = float_frame.loc[:, "A"]
+        tm.assert_series_equal(result, expected)
+
+        result = float_frame[lambda x: ["A", "B"]]
+        expected = float_frame.loc[:, ["A", "B"]]
+        tm.assert_frame_equal(result, float_frame.loc[:, ["A", "B"]])
+
+        df = float_frame[:3]
+        result = df[lambda x: [True, False, True]]
+        expected = float_frame.iloc[[0, 2], :]
+        tm.assert_frame_equal(result, expected)
+
+
+class TestGetitemBooleanMask:
+    def test_getitem_bool_mask_categorical_index(self):
+
+        df3 = DataFrame(
+            {
+                "A": np.arange(6, dtype="int64"),
+            },
+            index=CategoricalIndex(
+                [1, 1, 2, 1, 3, 2],
+                dtype=CategoricalDtype([3, 2, 1], ordered=True),
+                name="B",
+            ),
+        )
+        df4 = DataFrame(
+            {
+                "A": np.arange(6, dtype="int64"),
+            },
+            index=CategoricalIndex(
+                [1, 1, 2, 1, 3, 2],
+                dtype=CategoricalDtype([3, 2, 1], ordered=False),
+                name="B",
+            ),
+        )
+
+        result = df3[df3.index == "a"]
+        expected = df3.iloc[[]]
+        tm.assert_frame_equal(result, expected)
+
+        result = df4[df4.index == "a"]
+        expected = df4.iloc[[]]
+        tm.assert_frame_equal(result, expected)
+
+        result = df3[df3.index == 1]
+        expected = df3.iloc[[0, 1, 3]]
+        tm.assert_frame_equal(result, expected)
+
+        result = df4[df4.index == 1]
+        expected = df4.iloc[[0, 1, 3]]
+        tm.assert_frame_equal(result, expected)
+
+        # since we have an ordered categorical
+
+        # CategoricalIndex([1, 1, 2, 1, 3, 2],
+        #         categories=[3, 2, 1],
+        #         ordered=True,
+        #         name='B')
+        result = df3[df3.index < 2]
+        expected = df3.iloc[[4]]
+        tm.assert_frame_equal(result, expected)
+
+        result = df3[df3.index > 1]
+        expected = df3.iloc[[]]
+        tm.assert_frame_equal(result, expected)
+
+        # unordered
+        # cannot be compared
+
+        # CategoricalIndex([1, 1, 2, 1, 3, 2],
+        #         categories=[3, 2, 1],
+        #         ordered=False,
+        #         name='B')
+        msg = "Unordered Categoricals can only compare equality or not"
+        with pytest.raises(TypeError, match=msg):
+            df4[df4.index < 2]
+        with pytest.raises(TypeError, match=msg):
+            df4[df4.index > 1]

pandas/tests/frame/indexing/test_indexing.py

@@ -110,21 +110,6 @@ def test_getitem_listlike(self, idx_type, levels, float_frame):
         with pytest.raises(KeyError, match="not in index"):
             frame[idx]
 
-    def test_getitem_callable(self, float_frame):
-        # GH 12533
-        result = float_frame[lambda x: "A"]
-        expected = float_frame.loc[:, "A"]
-        tm.assert_series_equal(result, expected)
-
-        result = float_frame[lambda x: ["A", "B"]]
-        expected = float_frame.loc[:, ["A", "B"]]
-        tm.assert_frame_equal(result, float_frame.loc[:, ["A", "B"]])
-
-        df = float_frame[:3]
-        result = df[lambda x: [True, False, True]]
-        expected = float_frame.iloc[[0, 2], :]
-        tm.assert_frame_equal(result, expected)
-
     def test_setitem_list(self, float_frame):
 
         float_frame["E"] = "foo"

pandas/tests/indexing/test_at.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from pandas import DataFrame
+from pandas import DataFrame, Series
 import pandas._testing as tm
 
 
@@ -39,3 +39,72 @@ def test_at_with_duplicate_axes_requires_scalar_lookup(self):
             df.at[1, ["A"]] = 1
         with pytest.raises(ValueError, match=msg):
             df.at[:, "A"] = 1
+
+
+class TestAtErrors:
+    # TODO: De-duplicate/parametrize
+    #  test_at_series_raises_key_error, test_at_frame_raises_key_error,
+    #  test_at_series_raises_key_error2, test_at_frame_raises_key_error2
+
+    def test_at_series_raises_key_error(self):
+        # GH#31724 .at should match .loc
+
+        ser = Series([1, 2, 3], index=[3, 2, 1])
+        result = ser.at[1]
+        assert result == 3
+        result = ser.loc[1]
+        assert result == 3
+
+        with pytest.raises(KeyError, match="a"):
+            ser.at["a"]
+        with pytest.raises(KeyError, match="a"):
+            # .at should match .loc
+            ser.loc["a"]
+
+    def test_at_frame_raises_key_error(self):
+        # GH#31724 .at should match .loc
+
+        df = DataFrame({0: [1, 2, 3]}, index=[3, 2, 1])
+
+        result = df.at[1, 0]
+        assert result == 3
+        result = df.loc[1, 0]
+        assert result == 3
+
+        with pytest.raises(KeyError, match="a"):
+            df.at["a", 0]
+        with pytest.raises(KeyError, match="a"):
+            df.loc["a", 0]
+
+        with pytest.raises(KeyError, match="a"):
+            df.at[1, "a"]
+        with pytest.raises(KeyError, match="a"):
+            df.loc[1, "a"]
+
+    def test_at_series_raises_key_error2(self):
+        # at should not fallback
+        # GH#7814
+        # GH#31724 .at should match .loc
+        ser = Series([1, 2, 3], index=list("abc"))
+        result = ser.at["a"]
+        assert result == 1
+        result = ser.loc["a"]
+        assert result == 1
+
+        with pytest.raises(KeyError, match="^0$"):
+            ser.at[0]
+        with pytest.raises(KeyError, match="^0$"):
+            ser.loc[0]
+
+    def test_at_frame_raises_key_error2(self):
+        # GH#31724 .at should match .loc
+        df = DataFrame({"A": [1, 2, 3]}, index=list("abc"))
+        result = df.at["a", "A"]
+        assert result == 1
+        result = df.loc["a", "A"]
+        assert result == 1
+
+        with pytest.raises(KeyError, match="^0$"):
+            df.at["a", 0]
+        with pytest.raises(KeyError, match="^0$"):
+            df.loc["a", 0]

pandas/tests/indexing/test_categorical.py

@@ -2,7 +2,6 @@
 import pytest
 
 from pandas.core.dtypes.common import is_categorical_dtype
-from pandas.core.dtypes.dtypes import CategoricalDtype
 
 import pandas as pd
 from pandas import (
@@ -276,27 +275,6 @@ def test_slicing_doc_examples(self):
         )
         tm.assert_frame_equal(result, expected)
 
-    def test_getitem_category_type(self):
-        # GH 14580
-        # test iloc() on Series with Categorical data
-
-        s = Series([1, 2, 3]).astype("category")
-
-        # get slice
-        result = s.iloc[0:2]
-        expected = Series([1, 2]).astype(CategoricalDtype([1, 2, 3]))
-        tm.assert_series_equal(result, expected)
-
-        # get list of indexes
-        result = s.iloc[[0, 1]]
-        expected = Series([1, 2]).astype(CategoricalDtype([1, 2, 3]))
-        tm.assert_series_equal(result, expected)
-
-        # get boolean array
-        result = s.iloc[[True, False, False]]
-        expected = Series([1]).astype(CategoricalDtype([1, 2, 3]))
-        tm.assert_series_equal(result, expected)
-
     def test_loc_listlike(self):
 
         # list of labels
@@ -413,17 +391,6 @@ def test_loc_listlike_dtypes(self):
         with pytest.raises(KeyError, match=msg):
             df.loc[["a", "x"]]
 
-    def test_getitem_with_listlike(self):
-        # GH 16115
-        cats = Categorical([Timestamp("12-31-1999"), Timestamp("12-31-2000")])
-
-        expected = DataFrame(
-            [[1, 0], [0, 1]], dtype="uint8", index=[0, 1], columns=cats
-        )
-        dummies = pd.get_dummies(cats)
-        result = dummies[list(dummies.columns)]
-        tm.assert_frame_equal(result, expected)
-
     def test_ix_categorical_index(self):
         # GH 12531
         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):
         assert df.loc["B", 1] == 4
         assert df.at["B", 1] == 4
 
-    def test_getitem_bool_mask_categorical_index(self):
-
-        df3 = DataFrame(
-            {
-                "A": np.arange(6, dtype="int64"),
-            },
-            index=CategoricalIndex(
-                [1, 1, 2, 1, 3, 2], dtype=CDT([3, 2, 1], ordered=True), name="B"
-            ),
-        )
-        df4 = DataFrame(
-            {
-                "A": np.arange(6, dtype="int64"),
-            },
-            index=CategoricalIndex(
-                [1, 1, 2, 1, 3, 2], dtype=CDT([3, 2, 1], ordered=False), name="B"
-            ),
-        )
-
-        result = df3[df3.index == "a"]
-        expected = df3.iloc[[]]
-        tm.assert_frame_equal(result, expected)
-
-        result = df4[df4.index == "a"]
-        expected = df4.iloc[[]]
-        tm.assert_frame_equal(result, expected)
-
-        result = df3[df3.index == 1]
-        expected = df3.iloc[[0, 1, 3]]
-        tm.assert_frame_equal(result, expected)
-
-        result = df4[df4.index == 1]
-        expected = df4.iloc[[0, 1, 3]]
-        tm.assert_frame_equal(result, expected)
-
-        # since we have an ordered categorical
-
-        # CategoricalIndex([1, 1, 2, 1, 3, 2],
-        #         categories=[3, 2, 1],
-        #         ordered=True,
-        #         name='B')
-        result = df3[df3.index < 2]
-        expected = df3.iloc[[4]]
-        tm.assert_frame_equal(result, expected)
-
-        result = df3[df3.index > 1]
-        expected = df3.iloc[[]]
-        tm.assert_frame_equal(result, expected)
-
-        # unordered
-        # cannot be compared
-
-        # CategoricalIndex([1, 1, 2, 1, 3, 2],
-        #         categories=[3, 2, 1],
-        #         ordered=False,
-        #         name='B')
-        msg = "Unordered Categoricals can only compare equality or not"
-        with pytest.raises(TypeError, match=msg):
-            df4[df4.index < 2]
-        with pytest.raises(TypeError, match=msg):
-            df4[df4.index > 1]
-
     def test_indexing_with_category(self):
 
         # https://github.com/pandas-dev/pandas/issues/12564