TST/REF: collect xs tests

pandas-dev · jreback · Oct 29, 2020 · Oct 27, 2020 · Oct 27, 2020 · Oct 28, 2020
commit 9853d6357cf04f8bfc615f4c206a199f55f6a36e
diff --git a/pandas/tests/frame/indexing/test_xs.py b/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, 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,160 @@ 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"], [
10000
"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)