diff --git a/pandas/tests/generic/test_label_or_level_utils.py b/pandas/tests/generic/test_label_or_level_utils.py index 87bb6a58600f4..d1c85d770621c 100644 --- a/pandas/tests/generic/test_label_or_level_utils.py +++ b/pandas/tests/generic/test_label_or_level_utils.py @@ -69,11 +69,12 @@ def assert_level_reference(frame, levels, axis): # --------- def test_is_level_or_label_reference_df_simple(df_levels, axis): + axis = df_levels._get_axis_number(axis) # Compute expected labels and levels expected_labels, expected_levels = get_labels_levels(df_levels) # Transpose frame if axis == 1 - if axis in {1, "columns"}: + if axis == 1: df_levels = df_levels.T # Perform checks @@ -83,8 +84,10 @@ def test_is_level_or_label_reference_df_simple(df_levels, axis): def test_is_level_reference_df_ambig(df_ambig, axis): + axis = df_ambig._get_axis_number(axis) + # Transpose frame if axis == 1 - if axis in {1, "columns"}: + if axis == 1: df_ambig = df_ambig.T # df has both an on-axis level and off-axis label named L1 @@ -130,15 +133,14 @@ def test_is_level_reference_series_axis1_error(df): # --------- def test_check_label_or_level_ambiguity_df(df_ambig, axis): + axis = df_ambig._get_axis_number(axis) # Transpose frame if axis == 1 - if axis in {1, "columns"}: + if axis == 1: df_ambig = df_ambig.T - - if axis in {0, "index"}: - msg = "'L1' is both an index level and a column label" - else: msg = "'L1' is both a column level and an index label" + else: + msg = "'L1' is both an index level and a column label" # df_ambig has both an on-axis level and off-axis label named L1 # Therefore, L1 is ambiguous. with pytest.raises(ValueError, match=msg): @@ -181,8 +183,9 @@ def test_check_label_or_level_ambiguity_series_axis1_error(df): # Test _get_label_or_level_values # =============================== def assert_label_values(frame, labels, axis): + axis = frame._get_axis_number(axis) for label in labels: - if axis in {0, "index"}: + if axis == 0: expected = frame[label]._values else: expected = frame.loc[label]._values @@ -192,8 +195,9 @@ def assert_label_values(frame, labels, axis): def assert_level_values(frame, levels, axis): + axis = frame._get_axis_number(axis) for level in levels: - if axis in {0, "index"}: + if axis == 0: expected = frame.index.get_level_values(level=level)._values else: expected = frame.columns.get_level_values(level=level)._values @@ -209,8 +213,9 @@ def test_get_label_or_level_values_df_simple(df_levels, axis): # Compute expected labels and levels expected_labels, expected_levels = get_labels_levels(df_levels) + axis = df_levels._get_axis_number(axis) # Transpose frame if axis == 1 - if axis in {1, "columns"}: + if axis == 1: df_levels = df_levels.T # Perform checks @@ -220,8 +225,9 @@ def test_get_label_or_level_values_df_simple(df_levels, axis): def test_get_label_or_level_values_df_ambig(df_ambig, axis): + axis = df_ambig._get_axis_number(axis) # Transpose frame if axis == 1 - if axis in {1, "columns"}: + if axis == 1: df_ambig = df_ambig.T # df has an on-axis level named L2, and it is not ambiguous. @@ -233,8 +239,9 @@ def test_get_label_or_level_values_df_ambig(df_ambig, axis): def test_get_label_or_level_values_df_duplabels(df_duplabels, axis): + axis = df_duplabels._get_axis_number(axis) # Transpose frame if axis == 1 - if axis in {1, "columns"}: + if axis == 1: df_duplabels = df_duplabels.T # df has unambiguous level 'L1' @@ -244,7 +251,7 @@ def test_get_label_or_level_values_df_duplabels(df_duplabels, axis): assert_label_values(df_duplabels, ["L3"], axis=axis) # df has duplicate labels 'L2' - if axis in {0, "index"}: + if axis == 0: expected_msg = "The column label 'L2' is not unique" else: expected_msg = "The index label 'L2' is not unique" @@ -278,10 +285,11 @@ def test_get_label_or_level_values_series_axis1_error(df): # Test _drop_labels_or_levels # =========================== def assert_labels_dropped(frame, labels, axis): + axis = frame._get_axis_number(axis) for label in labels: df_dropped = frame._drop_labels_or_levels(label, axis=axis) - if axis in {0, "index"}: + if axis == 0: assert label in frame.columns assert label not in df_dropped.columns else: @@ -290,10 +298,11 @@ def assert_labels_dropped(frame, labels, axis): def assert_levels_dropped(frame, levels, axis): + axis = frame._get_axis_number(axis) for level in levels: df_dropped = frame._drop_labels_or_levels(level, axis=axis) - if axis in {0, "index"}: + if axis == 0: assert level in frame.index.names assert level not in df_dropped.index.names else: @@ -308,8 +317,9 @@ def test_drop_labels_or_levels_df(df_levels, axis): # Compute expected labels and levels expected_labels, expected_levels = get_labels_levels(df_levels) + axis = df_levels._get_axis_number(axis) # Transpose frame if axis == 1 - if axis in {1, "columns"}: + if axis == 1: df_levels = df_levels.T # Perform checks