TST: Add tests for single level indexing with loc(axis=1) (#29519)

Author: ryankarlos

pandas/tests/indexing/multiindex/test_slice.py

@@ -492,6 +492,44 @@ def test_loc_axis_arguments(self):
  with pytest.raises(ValueError):
  df.loc(axis="foo")[:, :, ["C1", "C3"]]
 
+ def test_loc_axis_single_level_multi_col_indexing_multiindex_col_df(self):
+
+ # GH29519
+ df = pd.DataFrame(
+ np.arange(27).reshape(3, 9),
+ columns=pd.MultiIndex.from_product(
+ [["a1", "a2", "a3"], ["b1", "b2", "b3"]]
+ ),
+ )
+ result = df.loc(axis=1)["a1":"a2"]
+ expected = df.iloc[:, :-3]
+
+ tm.assert_frame_equal(result, expected)
+
+ def test_loc_axis_single_level_single_col_indexing_multiindex_col_df(self):
+
+ # GH29519
+ df = pd.DataFrame(
+ np.arange(27).reshape(3, 9),
+ columns=pd.MultiIndex.from_product(
+ [["a1", "a2", "a3"], ["b1", "b2", "b3"]]
+ ),
+ )
+ result = df.loc(axis=1)["a1"]
+ expected = df.iloc[:, :3]
+ expected.columns = ["b1", "b2", "b3"]
+
+ tm.assert_frame_equal(result, expected)
+
+ def test_loc_ax_single_level_indexer_simple_df(self):
+
+ # GH29519
+ # test single level indexing on single index column data frame
+ df = pd.DataFrame(np.arange(9).reshape(3, 3), columns=["a", "b", "c"])
+ result = df.loc(axis=1)["a"]
+ expected = pd.Series(np.array([0, 3, 6]), name="a")
+ tm.assert_series_equal(result, expected)
+
  def test_per_axis_per_level_setitem(self):
 
  # test index maker