TST/REF: misplaced .xs tests (#37456)

Author: jbrockmendel

pandas/tests/frame/indexing/test_xs.py

@@ -3,12 +3,30 @@
 import numpy as np
 import pytest
 
-from pandas import DataFrame, Index, Series
+from pandas import DataFrame, Index, IndexSlice, MultiIndex, Series, concat
 import pandas._testing as tm
+import pandas.core.common as com
 
 from pandas.tseries.offsets import BDay
 
 
+@pytest.fixture
+def four_level_index_dataframe():
+    arr = np.array(
+        [
+            [-0.5109, -2.3358, -0.4645, 0.05076, 0.364],
+            [0.4473, 1.4152, 0.2834, 1.00661, 0.1744],
+            [-0.6662, -0.5243, -0.358, 0.89145, 2.5838],
+        ]
+    )
+    index = MultiIndex(
+        levels=[["a", "x"], ["b", "q"], [10.0032, 20.0, 30.0], [3, 4, 5]],
+        codes=[[0, 0, 1], [0, 1, 1], [0, 1, 2], [2, 1, 0]],
+        names=["one", "two", "three", "four"],
+    )
+    return DataFrame(arr, index=index, columns=list("ABCDE"))
+
+
 class TestXS:
     def test_xs(self, float_frame, datetime_frame):
         idx = float_frame.index[5]
@@ -92,3 +110,190 @@ def test_xs_view(self):
 
         dm.xs(2)[:] = 10
         assert (dm.xs(2) == 10).all()
+
+
+class TestXSWithMultiIndex:
+    def test_xs_integer_key(self):
+        # see GH#2107
+        dates = range(20111201, 20111205)
+        ids = list("abcde")
+        index = MultiIndex.from_product([dates, ids], names=["date", "secid"])
+        df = DataFrame(np.random.randn(len(index), 3), index, ["X", "Y", "Z"])
+
+        result = df.xs(20111201, level="date")
+        expected = df.loc[20111201, :]
+        tm.assert_frame_equal(result, expected)
+
+    def test_xs_level(self, multiindex_dataframe_random_data):
+        df = multiindex_dataframe_random_data
+        result = df.xs("two", level="second")
+        expected = df[df.index.get_level_values(1) == "two"]
+        expected.index = Index(["foo", "bar", "baz", "qux"], name="first")
+        tm.assert_frame_equal(result, expected)
+
+    def test_xs_level_eq_2(self):
+        arr = np.random.randn(3, 5)
+        index = MultiIndex(
+            levels=[["a", "p", "x"], ["b", "q", "y"], ["c", "r", "z"]],
+            codes=[[2, 0, 1], [2, 0, 1], [2, 0, 1]],
+        )
+        df = DataFrame(arr, index=index)
+        expected = DataFrame(arr[1:2], index=[["a"], ["b"]])
+        result = df.xs("c", level=2)
+        tm.assert_frame_equal(result, expected)
+
+    def test_xs_setting_with_copy_error(self, multiindex_dataframe_random_data):
+        # this is a copy in 0.14
+        df = multiindex_dataframe_random_data
+        result = df.xs("two", level="second")
+
+        # setting this will give a SettingWithCopyError
+        # as we are trying to write a view
+        msg = "A value is trying to be set on a copy of a slice from a DataFrame"
+        with pytest.raises(com.SettingWithCopyError, match=msg):
+            result[:] = 10
+
+    def test_xs_setting_with_copy_error_multiple(self, four_level_index_dataframe):
+        # this is a copy in 0.14
+        df = four_level_index_dataframe
+        result = df.xs(("a", 4), level=["one", "four"])
+
+        # setting this will give a SettingWithCopyError
+        # as we are trying to write a view
+        msg = "A value is trying to be set on a copy of a slice from a DataFrame"
+        with pytest.raises(com.SettingWithCopyError, match=msg):
+            result[:] = 10
+
+    @pytest.mark.parametrize("key, level", [("one", "second"), (["one"], ["second"])])
+    def test_xs_with_duplicates(self, key, level, multiindex_dataframe_random_data):
+        # see GH#13719
+        frame = multiindex_dataframe_random_data
+        df = concat([frame] * 2)
+        assert df.index.is_unique is False
+        expected = concat([frame.xs("one", level="second")] * 2)
+
+        result = df.xs(key, level=level)
+        tm.assert_frame_equal(result, expected)
+
+    def test_xs_missing_values_in_index(self):
+        # see GH#6574
+        # missing values in returned index should be preserved
+        acc = [
+            ("a", "abcde", 1),
+            ("b", "bbcde", 2),
+            ("y", "yzcde", 25),
+            ("z", "xbcde", 24),
+            ("z", None, 26),
+            ("z", "zbcde", 25),
+            ("z", "ybcde", 26),
+        ]
+        df = DataFrame(acc, columns=["a1", "a2", "cnt"]).set_index(["a1", "a2"])
+        expected = DataFrame(
+            {"cnt": [24, 26, 25, 26]},
+            index=Index(["xbcde", np.nan, "zbcde", "ybcde"], name="a2"),
+        )
+
+        result = df.xs("z", level="a1")
+        tm.assert_frame_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "key, level, exp_arr, exp_index",
+        [
+            ("a", "lvl0", lambda x: x[:, 0:2], Index(["bar", "foo"], name="lvl1")),
+            ("foo", "lvl1", lambda x: x[:, 1:2], Index(["a"], name="lvl0")),
+        ],
+    )
+    def test_xs_named_levels_axis_eq_1(self, key, level, exp_arr, exp_index):
+        # see GH#2903
+        arr = np.random.randn(4, 4)
+        index = MultiIndex(
+            levels=[["a", "b"], ["bar", "foo", "hello", "world"]],
+            codes=[[0, 0, 1, 1], [0, 1, 2, 3]],
+            names=["lvl0", "lvl1"],
+        )
+        df = DataFrame(arr, columns=index)
+        result = df.xs(key, level=level, axis=1)
+        expected = DataFrame(exp_arr(arr), columns=exp_index)
+        tm.assert_frame_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "indexer",
+        [
+            lambda df: df.xs(("a", 4), level=["one", "four"]),
+            lambda df: df.xs("a").xs(4, level="four"),
+        ],
+    )
+    def test_xs_level_multiple(self, indexer, four_level_index_dataframe):
+        df = four_level_index_dataframe
+        expected_values = [[0.4473, 1.4152, 0.2834, 1.00661, 0.1744]]
+        expected_index = MultiIndex(
+            levels=[["q"], [20.0]], codes=[[0], [0]], names=["two", "three"]
+        )
+        expected = DataFrame(
+            expected_values, index=expected_index, columns=list("ABCDE")
+        )
+        result = indexer(df)
+        tm.assert_frame_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "indexer", [lambda df: df.xs("a", level=0), lambda df: df.xs("a")]
+    )
+    def test_xs_level0(self, indexer, four_level_index_dataframe):
+        df = four_level_index_dataframe
+        expected_values = [
+            [-0.5109, -2.3358, -0.4645, 0.05076, 0.364],
+            [0.4473, 1.4152, 0.2834, 1.00661, 0.1744],
+        ]
+        expected_index = MultiIndex(
+            levels=[["b", "q"], [10.0032, 20.0], [4, 5]],
+            codes=[[0, 1], [0, 1], [1, 0]],
+            names=["two", "three", "four"],
+        )
+        expected = DataFrame(
+            expected_values, index=expected_index, columns=list("ABCDE")
+        )
+
+        result = indexer(df)
+        tm.assert_frame_equal(result, expected)
+
+    def test_xs_values(self, multiindex_dataframe_random_data):
+        df = multiindex_dataframe_random_data
+        result = df.xs(("bar", "two")).values
+        expected = df.values[4]
+        tm.assert_almost_equal(result, expected)
+
+    def test_xs_loc_equality(self, multiindex_dataframe_random_data):
+        df = multiindex_dataframe_random_data
+        result = df.xs(("bar", "two"))
+        expected = df.loc[("bar", "two")]
+        tm.assert_series_equal(result, expected)
+
+    @pytest.mark.parametrize("klass", [DataFrame, Series])
+    def test_xs_IndexSlice_argument_not_implemented(self, klass):
+        # GH#35301
+
+        index = MultiIndex(
+            levels=[[("foo", "bar", 0), ("foo", "baz", 0), ("foo", "qux", 0)], [0, 1]],
+            codes=[[0, 0, 1, 1, 2, 2], [0, 1, 0, 1, 0, 1]],
+        )
+
+        obj = DataFrame(np.random.randn(6, 4), index=index)
+        if klass is Series:
+            obj = obj[0]
+
+        msg = (
+            "Expected label or tuple of labels, got "
+            r"\(\('foo', 'qux', 0\), slice\(None, None, None\)\)"
+        )
+        with pytest.raises(TypeError, match=msg):
+            obj.xs(IndexSlice[("foo", "qux", 0), :])
+
+    @pytest.mark.parametrize("klass", [DataFrame, Series])
+    def test_xs_levels_raises(self, klass):
+        obj = DataFrame({"A": [1, 2, 3]})
+        if klass is Series:
+            obj = obj["A"]
+
+        msg = "Index must be a MultiIndex"
+        with pytest.raises(TypeError, match=msg):
+            obj.xs(0, level="as")

pandas/tests/indexes/multi/test_indexing.py

@@ -662,6 +662,13 @@ def test_get_loc_past_lexsort_depth(self):
 
         assert result == slice(0, 1, None)
 
+    def test_multiindex_get_loc_list_raises(self):
+        # GH#35878
+        idx = MultiIndex.from_tuples([("a", 1), ("b", 2)])
+        msg = "unhashable type"
+        with pytest.raises(TypeError, match=msg):
+            idx.get_loc([])
+
 
 class TestWhere:
     def test_where(self):

pandas/tests/indexing/multiindex/test_loc.py

@@ -590,3 +590,11 @@ def test_missing_key_raises_keyerror2(self):
 
         with pytest.raises(KeyError, match=r"\(0, 3\)"):
             ser.loc[0, 3]
+
+
+def test_getitem_loc_commutability(multiindex_year_month_day_dataframe_random_data):
+    df = multiindex_year_month_day_dataframe_random_data
+    ser = df["A"]
+    result = ser[2000, 5]
+    expected = df.loc[2000, 5]["A"]
+    tm.assert_series_equal(result, expected)

pandas/tests/indexing/multiindex/test_multiindex.py

@@ -1,5 +1,4 @@
 import numpy as np
-import pytest
 
 import pandas._libs.index as _index
 from pandas.errors import PerformanceWarning
@@ -84,10 +83,3 @@ def test_nested_tuples_duplicates(self):
         df3 = df.copy(deep=True)
         df3.loc[[(dti[0], "a")], "c2"] = 1.0
         tm.assert_frame_equal(df3, expected)
-
-    def test_multiindex_get_loc_list_raises(self):
-        # https://github.com/pandas-dev/pandas/issues/35878
-        idx = pd.MultiIndex.from_tuples([("a", 1), ("b", 2)])
-        msg = "unhashable type"
-        with pytest.raises(TypeError, match=msg):
-            idx.get_loc([])