REF: Replacing axis in with _get_axis_number(axis) (#43131)

Author: RohanTrix

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