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

Author: jbrockmendel

pandas/tests/frame/indexing/test_getitem.py

@@ -68,6 +68,18 @@ def test_getitem_sparse_column_return_type_and_dtype(self):
         tm.assert_series_equal(result, expected)
 
 
+class TestGetitemListLike:
+    def test_getitem_list_missing_key(self):
+        # GH#13822, incorrect error string with non-unique columns when missing
+        # column is accessed
+        df = DataFrame({"x": [1.0], "y": [2.0], "z": [3.0]})
+        df.columns = ["x", "x", "z"]
+
+        # Check that we get the correct value in the KeyError
+        with pytest.raises(KeyError, match=r"\['y'\] not in index"):
+            df[["x", "y", "z"]]
+
+
 class TestGetitemCallable:
     def test_getitem_callable(self, float_frame):
         # GH#12533

pandas/tests/indexing/test_at.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from pandas import DataFrame, Series
+from pandas import DataFrame, Series, Timestamp
 import pandas._testing as tm
 
 
@@ -27,6 +27,16 @@ def test_at_setitem_mixed_index_assignment(self):
         assert ser.iat[3] == 22
 
 
+class TestAtSetItemWithExpansion:
+    def test_at_setitem_expansion_series_dt64tz_value(self, tz_naive_fixture):
+        # GH#25506
+        ts = Timestamp("2017-08-05 00:00:00+0100", tz=tz_naive_fixture)
+        result = Series(ts)
+        result.at[1] = ts
+        expected = Series([ts, ts])
+        tm.assert_series_equal(result, expected)
+
+
 class TestAtWithDuplicates:
     def test_at_with_duplicate_axes_requires_scalar_lookup(self):
         # GH#33041 check that falling back to loc doesn't allow non-scalar

pandas/tests/indexing/test_categorical.py

@@ -430,25 +430,6 @@ def test_ix_categorical_index(self):
         )
         tm.assert_frame_equal(cdf.loc[:, ["X", "Y"]], expect)
 
-    def test_read_only_source(self):
-        # GH 10043
-        rw_array = np.eye(10)
-        rw_df = DataFrame(rw_array)
-
-        ro_array = np.eye(10)
-        ro_array.setflags(write=False)
-        ro_df = DataFrame(ro_array)
-
-        tm.assert_frame_equal(rw_df.iloc[[1, 2, 3]], ro_df.iloc[[1, 2, 3]])
-        tm.assert_frame_equal(rw_df.iloc[[1]], ro_df.iloc[[1]])
-        tm.assert_series_equal(rw_df.iloc[1], ro_df.iloc[1])
-        tm.assert_frame_equal(rw_df.iloc[1:3], ro_df.iloc[1:3])
-
-        tm.assert_frame_equal(rw_df.loc[[1, 2, 3]], ro_df.loc[[1, 2, 3]])
-        tm.assert_frame_equal(rw_df.loc[[1]], ro_df.loc[[1]])
-        tm.assert_series_equal(rw_df.loc[1], ro_df.loc[1])
-        tm.assert_frame_equal(rw_df.loc[1:3], ro_df.loc[1:3])
-
     def test_loc_slice(self):
         # GH9748
         with pytest.raises(KeyError, match="1"):

pandas/tests/indexing/test_datetime.py

@@ -160,24 +160,33 @@ def test_indexing_with_datetimeindex_tz(self):
         expected = Series([0, 5], index=index)
         tm.assert_series_equal(result, expected)
 
-    def test_series_partial_set_datetime(self):
+    @pytest.mark.parametrize("to_period", [True, False])
+    def test_loc_getitem_listlike_of_datetimelike_keys(self, to_period):
         # GH 11497
 
         idx = date_range("2011-01-01", "2011-01-02", freq="D", name="idx")
+        if to_period:
+            idx = idx.to_period("D")
         ser = Series([0.1, 0.2], index=idx, name="s")
 
-        result = ser.loc[[Timestamp("2011-01-01"), Timestamp("2011-01-02")]]
+        keys = [Timestamp("2011-01-01"), Timestamp("2011-01-02")]
+        if to_period:
+            keys = [x.to_period("D") for x in keys]
+        result = ser.loc[keys]
         exp = Series([0.1, 0.2], index=idx, name="s")
-        exp.index = exp.index._with_freq(None)
+        if not to_period:
+            exp.index = exp.index._with_freq(None)
         tm.assert_series_equal(result, exp, check_index_type=True)
 
         keys = [
             Timestamp("2011-01-02"),
             Timestamp("2011-01-02"),
             Timestamp("2011-01-01"),
         ]
+        if to_period:
+            keys = [x.to_period("D") for x in keys]
         exp = Series(
-            [0.2, 0.2, 0.1], index=pd.DatetimeIndex(keys, name="idx"), name="s"
+            [0.2, 0.2, 0.1], index=Index(keys, name="idx", dtype=idx.dtype), name="s"
         )
         result = ser.loc[keys]
         tm.assert_series_equal(result, exp, check_index_type=True)
@@ -187,35 +196,9 @@ def test_series_partial_set_datetime(self):
             Timestamp("2011-01-02"),
             Timestamp("2011-01-03"),
         ]
-        with pytest.raises(KeyError, match="with any missing labels"):
-            ser.loc[keys]
-
-    def test_series_partial_set_period(self):
-        # GH 11497
-
-        idx = pd.period_range("2011-01-01", "2011-01-02", freq="D", name="idx")
-        ser = Series([0.1, 0.2], index=idx, name="s")
-
-        result = ser.loc[
-            [pd.Period("2011-01-01", freq="D"), pd.Period("2011-01-02", freq="D")]
-        ]
-        exp = Series([0.1, 0.2], index=idx, name="s")
-        tm.assert_series_equal(result, exp, check_index_type=True)
+        if to_period:
+            keys = [x.to_period("D") for x in keys]
 
-        keys = [
-            pd.Period("2011-01-02", freq="D"),
-            pd.Period("2011-01-02", freq="D"),
-            pd.Period("2011-01-01", freq="D"),
-        ]
-        exp = Series([0.2, 0.2, 0.1], index=pd.PeriodIndex(keys, name="idx"), name="s")
-        result = ser.loc[keys]
-        tm.assert_series_equal(result, exp, check_index_type=True)
-
-        keys = [
-            pd.Period("2011-01-03", freq="D"),
-            pd.Period("2011-01-02", freq="D"),
-            pd.Period("2011-01-03", freq="D"),
-        ]
         with pytest.raises(KeyError, match="with any missing labels"):
             ser.loc[keys]
 

pandas/tests/indexing/test_floats.py

@@ -140,8 +140,11 @@ def test_scalar_with_mixed(self):
         expected = 3
         assert result == expected
 
+    @pytest.mark.parametrize(
+        "idxr,getitem", [(lambda x: x.loc, False), (lambda x: x, True)]
+    )
     @pytest.mark.parametrize("index_func", [tm.makeIntIndex, tm.makeRangeIndex])
-    def test_scalar_integer(self, index_func, frame_or_series):
+    def test_scalar_integer(self, index_func, frame_or_series, idxr, getitem):
 
         # test how scalar float indexers work on int indexes
 
@@ -150,37 +153,39 @@ def test_scalar_integer(self, index_func, frame_or_series):
         obj = gen_obj(frame_or_series, i)
 
         # coerce to equal int
-        for idxr, getitem in [(lambda x: x.loc, False), (lambda x: x, True)]:
 
-            result = idxr(obj)[3.0]
-            self.check(result, obj, 3, getitem)
+        result = idxr(obj)[3.0]
+        self.check(result, obj, 3, getitem)
 
-        # coerce to equal int
-        for idxr, getitem in [(lambda x: x.loc, False), (lambda x: x, True)]:
-
-            if isinstance(obj, Series):
+        if isinstance(obj, Series):
 
-                def compare(x, y):
-                    assert x == y
+            def compare(x, y):
+                assert x == y
 
-                expected = 100
+            expected = 100
+        else:
+            compare = tm.assert_series_equal
+            if getitem:
+                expected = Series(100, index=range(len(obj)), name=3)
             else:
-                compare = tm.assert_series_equal
-                if getitem:
-                    expected = Series(100, index=range(len(obj)), name=3)
-                else:
-                    expected = Series(100.0, index=range(len(obj)), name=3)
+                expected = Series(100.0, index=range(len(obj)), name=3)
 
-            s2 = obj.copy()
-            idxr(s2)[3.0] = 100
+        s2 = obj.copy()
+        idxr(s2)[3.0] = 100
 
-            result = idxr(s2)[3.0]
-            compare(result, expected)
+        result = idxr(s2)[3.0]
+        compare(result, expected)
 
-            result = idxr(s2)[3]
-            compare(result, expected)
+        result = idxr(s2)[3]
+        compare(result, expected)
 
+    @pytest.mark.parametrize("index_func", [tm.makeIntIndex, tm.makeRangeIndex])
+    def test_scalar_integer_contains_float(self, index_func, frame_or_series):
         # contains
+        # integer index
+        index = index_func(5)
+        obj = gen_obj(frame_or_series, index)
+
         # coerce to equal int
         assert 3.0 in obj
 

pandas/tests/indexing/test_iat.py

@@ -1,4 +1,6 @@
-import pandas as pd
+import numpy as np
+
+from pandas import DataFrame, Series, period_range
 
 
 def test_iat(float_frame):
@@ -12,5 +14,14 @@ def test_iat(float_frame):
 
 def test_iat_duplicate_columns():
     # https://github.com/pandas-dev/pandas/issues/11754
-    df = pd.DataFrame([[1, 2]], columns=["x", "x"])
+    df = DataFrame([[1, 2]], columns=["x", "x"])
     assert df.iat[0, 0] == 1
+
+
+def test_iat_getitem_series_with_period_index():
+    # GH#4390, iat incorrectly indexing
+    index = period_range("1/1/2001", periods=10)
+    ser = Series(np.random.randn(10), index=index)
+    expected = ser[index[0]]
+    result = ser.iat[0]
+    assert expected == result

pandas/tests/indexing/test_iloc.py

@@ -801,6 +801,36 @@ def test_iloc_setitem_empty_frame_raises_with_3d_ndarray(self):
         with pytest.raises(ValueError, match=msg):
             obj.iloc[nd3] = 0
 
+    @pytest.mark.parametrize("indexer", [lambda x: x.loc, lambda x: x.iloc])
+    def test_iloc_getitem_read_only_values(self, indexer):
+        # GH#10043 this is fundamentally a test for iloc, but test loc while
+        #  we're here
+        rw_array = np.eye(10)
+        rw_df = DataFrame(rw_array)
+
+        ro_array = np.eye(10)
+        ro_array.setflags(write=False)
+        ro_df = DataFrame(ro_array)
+
+        tm.assert_frame_equal(indexer(rw_df)[[1, 2, 3]], indexer(ro_df)[[1, 2, 3]])
+        tm.assert_frame_equal(indexer(rw_df)[[1]], indexer(ro_df)[[1]])
+        tm.assert_series_equal(indexer(rw_df)[1], indexer(ro_df)[1])
+        tm.assert_frame_equal(indexer(rw_df)[1:3], indexer(ro_df)[1:3])
+
+    def test_iloc_getitem_readonly_key(self):
+        # GH#17192 iloc with read-only array raising TypeError
+        df = DataFrame({"data": np.ones(100, dtype="float64")})
+        indices = np.array([1, 3, 6])
+        indices.flags.writeable = False
+
+        result = df.iloc[indices]
+        expected = df.loc[[1, 3, 6]]
+        tm.assert_frame_equal(result, expected)
+
+        result = df["data"].iloc[indices]
+        expected = df["data"].loc[[1, 3, 6]]
+        tm.assert_series_equal(result, expected)
+
     def test_iloc_assign_series_to_df_cell(self):
         # GH 37593
         df = DataFrame(columns=["a"], index=[0])