TST: Split/parameterize indexing tests (#45325)

Author: mroeschke

pandas/conftest.py

@@ -1751,6 +1751,14 @@ def indexer_al(request):
     return request.param
 
 
+@pytest.fixture(params=[tm.iat, tm.iloc])
+def indexer_ial(request):
+    """
+    Parametrize over iat.__setitem__, iloc.__setitem__
+    """
+    return request.param
+
+
 @pytest.fixture
 def using_array_manager(request):
     """

pandas/tests/indexing/test_indexing.py

@@ -704,20 +704,20 @@ def run_tests(df, rhs, right_loc, right_iloc):
         right_iloc["jolie"] = ["@2", -26.0, -18.0, -10.0, "@18"]
         run_tests(df, rhs, right_loc, right_iloc)
 
-    def test_str_label_slicing_with_negative_step(self):
+    @pytest.mark.parametrize(
+        "idx", [_mklbl("A", 20), np.arange(20) + 100, np.linspace(100, 150, 20)]
+    )
+    def test_str_label_slicing_with_negative_step(self, idx):
         SLC = pd.IndexSlice
 
-        for idx in [_mklbl("A", 20), np.arange(20) + 100, np.linspace(100, 150, 20)]:
-            idx = Index(idx)
-            ser = Series(np.arange(20), index=idx)
-            tm.assert_indexing_slices_equivalent(ser, SLC[idx[9] :: -1], SLC[9::-1])
-            tm.assert_indexing_slices_equivalent(ser, SLC[: idx[9] : -1], SLC[:8:-1])
-            tm.assert_indexing_slices_equivalent(
-                ser, SLC[idx[13] : idx[9] : -1], SLC[13:8:-1]
-            )
-            tm.assert_indexing_slices_equivalent(
-                ser, SLC[idx[9] : idx[13] : -1], SLC[:0]
-            )
+        idx = Index(idx)
+        ser = Series(np.arange(20), index=idx)
+        tm.assert_indexing_slices_equivalent(ser, SLC[idx[9] :: -1], SLC[9::-1])
+        tm.assert_indexing_slices_equivalent(ser, SLC[: idx[9] : -1], SLC[:8:-1])
+        tm.assert_indexing_slices_equivalent(
+            ser, SLC[idx[13] : idx[9] : -1], SLC[13:8:-1]
+        )
+        tm.assert_indexing_slices_equivalent(ser, SLC[idx[9] : idx[13] : -1], SLC[:0])
 
     def test_slice_with_zero_step_raises(self, index, indexer_sl, frame_or_series):
         obj = frame_or_series(np.arange(len(index)), index=index)

pandas/tests/indexing/test_loc.py

@@ -54,42 +54,45 @@ def test_loc_getitem_label(self):
         # label
         self.check_result("loc", "c", typs=["empty"], fails=KeyError)
 
-    def test_loc_getitem_label_out_of_range(self):
+    @pytest.mark.parametrize(
+        "key, typs, axes",
+        [
+            ["f", ["ints", "uints", "labels", "mixed", "ts"], None],
+            ["f", ["floats"], None],
+            [20, ["ints", "uints", "mixed"], None],
+            [20, ["labels"], None],
+            [20, ["ts"], 0],
+            [20, ["floats"], 0],
+        ],
+    )
+    def test_loc_getitem_label_out_of_range(self, key, typs, axes):
 
         # out of range label
-        self.check_result(
-            "loc", "f", typs=["ints", "uints", "labels", "mixed", "ts"], fails=KeyError
-        )
-        self.check_result("loc", "f", typs=["floats"], fails=KeyError)
-        self.check_result("loc", "f", typs=["floats"], fails=KeyError)
-        self.check_result("loc", 20, typs=["ints", "uints", "mixed"], fails=KeyError)
-        self.check_result("loc", 20, typs=["labels"], fails=KeyError)
-        self.check_result("loc", 20, typs=["ts"], axes=0, fails=KeyError)
-        self.check_result("loc", 20, typs=["floats"], axes=0, fails=KeyError)
+        self.check_result("loc", key, typs=typs, axes=axes, fails=KeyError)
 
-    def test_loc_getitem_label_list(self):
+    @pytest.mark.parametrize(
+        "key, typs",
+        [
+            [[0, 1, 2], ["ints", "uints", "floats"]],
+            [[1, 3.0, "A"], ["ints", "uints", "floats"]],
+        ],
+    )
+    def test_loc_getitem_label_list(self, key, typs):
         # list of labels
-        self.check_result(
-            "loc", [0, 1, 2], typs=["ints", "uints", "floats"], fails=KeyError
-        )
-        self.check_result(
-            "loc", [1, 3.0, "A"], typs=["ints", "uints", "floats"], fails=KeyError
-        )
-
-    def test_loc_getitem_label_list_with_missing(self):
-        self.check_result("loc", [0, 1, 2], typs=["empty"], fails=KeyError)
-        self.check_result(
-            "loc", [0, 2, 10], typs=["ints", "uints", "floats"], axes=0, fails=KeyError
-        )
+        self.check_result("loc", key, typs=typs, fails=KeyError)
 
-        self.check_result(
-            "loc", [3, 6, 7], typs=["ints", "uints", "floats"], axes=1, fails=KeyError
-        )
-
-        # GH 17758 - MultiIndex and missing keys
-        self.check_result(
-            "loc", [(1, 3), (1, 4), (2, 5)], typs=["multi"], axes=0, fails=KeyError
-        )
+    @pytest.mark.parametrize(
+        "key, typs, axes",
+        [
+            [[0, 1, 2], ["empty"], None],
+            [[0, 2, 10], ["ints", "uints", "floats"], 0],
+            [[3, 6, 7], ["ints", "uints", "floats"], 1],
+            # GH 17758 - MultiIndex and missing keys
+            [[(1, 3), (1, 4), (2, 5)], ["multi"], 0],
+        ],
+    )
+    def test_loc_getitem_label_list_with_missing(self, key, typs, axes):
+        self.check_result("loc", key, typs=typs, axes=axes, fails=KeyError)
 
     def test_loc_getitem_label_list_fails(self):
         # fails
@@ -108,7 +111,22 @@ def test_loc_getitem_bool(self):
 
         self.check_result("loc", b, typs=["empty"], fails=IndexError)
 
-    def test_loc_getitem_label_slice(self):
+    @pytest.mark.parametrize(
+        "slc, typs, axes, fails",
+        [
+            [
+                slice(1, 3),
+                ["labels", "mixed", "empty", "ts", "floats"],
+                None,
+                TypeError,
+            ],
+            [slice("20130102", "20130104"), ["ts"], 1, TypeError],
+            [slice(2, 8), ["mixed"], 0, TypeError],
+            [slice(2, 8), ["mixed"], 1, KeyError],
+            [slice(2, 4, 2), ["mixed"], 0, TypeError],
+        ],
+    )
+    def test_loc_getitem_label_slice(self, slc, typs, axes, fails):
 
         # label slices (with ints)
 
@@ -118,20 +136,10 @@ def test_loc_getitem_label_slice(self):
 
         self.check_result(
             "loc",
-            slice(1, 3),
-            typs=["labels", "mixed", "empty", "ts", "floats"],
-            fails=TypeError,
-        )
-
-        self.check_result(
-            "loc", slice("20130102", "20130104"), typs=["ts"], axes=1, fails=TypeError
-        )
-
-        self.check_result("loc", slice(2, 8), typs=["mixed"], axes=0, fails=TypeError)
-        self.check_result("loc", slice(2, 8), typs=["mixed"], axes=1, fails=KeyError)
-
-        self.check_result(
-            "loc", slice(2, 4, 2), typs=["mixed"], axes=0, fails=TypeError
+            slc,
+            typs=typs,
+            axes=axes,
+            fails=fails,
         )
 
     def test_setitem_from_duplicate_axis(self):
@@ -956,55 +964,40 @@ def test_loc_non_unique(self):
         tm.assert_frame_equal(result, expected)
 
     @pytest.mark.arm_slow
-    def test_loc_non_unique_memory_error(self):
+    @pytest.mark.parametrize("length, l2", [[900, 100], [900000, 100000]])
+    def test_loc_non_unique_memory_error(self, length, l2):
 
         # GH 4280
         # non_unique index with a large selection triggers a memory error
 
         columns = list("ABCDEFG")
 
-        def gen_test(length, l2):
-            return pd.concat(
-                [
-                    DataFrame(
-                        np.random.randn(length, len(columns)),
-                        index=np.arange(length),
-                        columns=columns,
-                    ),
-                    DataFrame(
-                        np.ones((l2, len(columns))), index=[0] * l2, columns=columns
-                    ),
-                ]
-            )
-
-        def gen_expected(df, mask):
-            len_mask = len(mask)
-            return pd.concat(
-                [
-                    df.take([0]),
-                    DataFrame(
-                        np.ones((len_mask, len(columns))),
-                        index=[0] * len_mask,
-                        columns=columns,
-                    ),
-                    df.take(mask[1:]),
-                ]
-            )
-
-        df = gen_test(900, 100)
-        assert df.index.is_unique is False
-
-        mask = np.arange(100)
-        result = df.loc[mask]
-        expected = gen_expected(df, mask)
-        tm.assert_frame_equal(result, expected)
+        df = pd.concat(
+            [
+                DataFrame(
+                    np.random.randn(length, len(columns)),
+                    index=np.arange(length),
+                    columns=columns,
+                ),
+                DataFrame(np.ones((l2, len(columns))), index=[0] * l2, columns=columns),
+            ]
+        )
 
-        df = gen_test(900000, 100000)
         assert df.index.is_unique is False
 
-        mask = np.arange(100000)
+        mask = np.arange(l2)
         result = df.loc[mask]
-        expected = gen_expected(df, mask)
+        expected = pd.concat(
+            [
+                df.take([0]),
+                DataFrame(
+                    np.ones((len(mask), len(columns))),
+                    index=[0] * len(mask),
+                    columns=columns,
+                ),
+                df.take(mask[1:]),
+            ]
+        )
         tm.assert_frame_equal(result, expected)
 
     def test_loc_name(self):
@@ -1840,12 +1833,20 @@ def test_loc_setitem_empty_series(self):
         ser.loc[3] = 3
         tm.assert_series_equal(ser, Series([1, 3], index=[1, 3]))
 
+    def test_loc_setitem_empty_series_float(self):
+        # GH#5226
+
+        # partially set with an empty object series
         ser = Series(dtype=object)
         ser.loc[1] = 1.0
         tm.assert_series_equal(ser, Series([1.0], index=[1]))
         ser.loc[3] = 3.0
         tm.assert_series_equal(ser, Series([1.0, 3.0], index=[1, 3]))
 
+    def test_loc_setitem_empty_series_str_idx(self):
+        # GH#5226
+
+        # partially set with an empty object series
         ser = Series(dtype=object)
         ser.loc["foo"] = 1
         tm.assert_series_equal(ser, Series([1], index=["foo"]))
@@ -1864,24 +1865,26 @@ def test_loc_setitem_incremental_with_dst(self):
         expected = Series(1, index=idxs)
         tm.assert_series_equal(result, expected)
 
-    def test_loc_setitem_datetime_keys_cast(self):
-        # GH#9516
-        dt1 = Timestamp("20130101 09:00:00")
-        dt2 = Timestamp("20130101 10:00:00")
-
-        for conv in [
+    @pytest.mark.parametrize(
+        "conv",
+        [
             lambda x: x,
             lambda x: x.to_datetime64(),
             lambda x: x.to_pydatetime(),
             lambda x: np.datetime64(x),
-        ]:
-
-            df = DataFrame()
-            df.loc[conv(dt1), "one"] = 100
-            df.loc[conv(dt2), "one"] = 200
+        ],
+        ids=["self", "to_datetime64", "to_pydatetime", "np.datetime64"],
+    )
+    def test_loc_setitem_datetime_keys_cast(self, conv):
+        # GH#9516
+        dt1 = Timestamp("20130101 09:00:00")
+        dt2 = Timestamp("20130101 10:00:00")
+        df = DataFrame()
+        df.loc[conv(dt1), "one"] = 100
+        df.loc[conv(dt2), "one"] = 200
 
-            expected = DataFrame({"one": [100.0, 200.0]}, index=[dt1, dt2])
-            tm.assert_frame_equal(df, expected)
+        expected = DataFrame({"one": [100.0, 200.0]}, index=[dt1, dt2])
+        tm.assert_frame_equal(df, expected)
 
     def test_loc_setitem_categorical_column_retains_dtype(self, ordered):
         # GH16360

pandas/tests/indexing/test_scalar.py

@@ -19,33 +19,6 @@
 
 
 class TestScalar(Base):
-    @pytest.mark.parametrize("kind", ["series", "frame"])
-    def test_at_and_iat_get(self, kind):
-        def _check(f, func, values=False):
-
-            if f is not None:
-                indices = self.generate_indices(f, values)
-                for i in indices:
-                    result = getattr(f, func)[i]
-                    expected = self.get_value(func, f, i, values)
-                    tm.assert_almost_equal(result, expected)
-
-        d = getattr(self, kind)
-
-        # iat
-        for f in [d["ints"], d["uints"]]:
-            _check(f, "iat", values=True)
-
-        for f in [d["labels"], d["ts"], d["floats"]]:
-            if f is not None:
-                msg = "iAt based indexing can only have integer indexers"
-                with pytest.raises(ValueError, match=msg):
-                    self.check_values(f, "iat")
-
-        # at
-        for f in [d["ints"], d["uints"], d["labels"], d["ts"], d["floats"]]:
-            _check(f, "at")
-
     @pytest.mark.parametrize("kind", ["series", "frame"])
     @pytest.mark.parametrize("col", ["ints", "uints"])
     def test_iat_set_ints(self, kind, col):
@@ -103,24 +76,24 @@ def test_at_iat_coercion(self):
         xp = s.values[5]
         assert result == xp
 
+    @pytest.mark.parametrize(
+        "ser, expected",
+        [
+            [
+                Series(["2014-01-01", "2014-02-02"], dtype="datetime64[ns]"),
+                Timestamp("2014-02-02"),
+            ],
+            [
+                Series(["1 days", "2 days"], dtype="timedelta64[ns]"),
+                Timedelta("2 days"),
+            ],
+        ],
+    )
+    def test_iloc_iat_coercion_datelike(self, indexer_ial, ser, expected):
         # GH 7729
         # make sure we are boxing the returns
-        s = Series(["2014-01-01", "2014-02-02"], dtype="datetime64[ns]")
-        expected = Timestamp("2014-02-02")
-
-        for r in [lambda: s.iat[1], lambda: s.iloc[1]]:
-            result = r()
-            assert result == expected
-
-        s = Series(["1 days", "2 days"], dtype="timedelta64[ns]")
-        expected = Timedelta("2 days")
-
-        for r in [lambda: s.iat[1], lambda: s.iloc[1]]:
-            result = r()
-            assert result == expected
-
-    def test_iat_invalid_args(self):
-        pass
+        result = indexer_ial(ser)[1]
+        assert result == expected
 
     def test_imethods_with_dups(self):