REF: set_axis tests (#33189)

Author: jbrockmendel

pandas/tests/frame/test_alter_axes.py

@@ -776,35 +776,3 @@ def test_set_reset_index(self):
         df = df.set_index("B")
 
         df = df.reset_index()
-
-    def test_set_axis_inplace(self):
-        # GH14636
-        df = DataFrame(
-            {"A": [1.1, 2.2, 3.3], "B": [5.0, 6.1, 7.2], "C": [4.4, 5.5, 6.6]},
-            index=[2010, 2011, 2012],
-        )
-
-        expected = {0: df.copy(), 1: df.copy()}
-        expected[0].index = list("abc")
-        expected[1].columns = list("abc")
-        expected["index"] = expected[0]
-        expected["columns"] = expected[1]
-
-        for axis in expected:
-            result = df.copy()
-            result.set_axis(list("abc"), axis=axis, inplace=True)
-            tm.assert_frame_equal(result, expected[axis])
-
-            # inplace=False
-            result = df.set_axis(list("abc"), axis=axis)
-            tm.assert_frame_equal(expected[axis], result)
-
-        # omitting the "axis" parameter
-        with tm.assert_produces_warning(None):
-            result = df.set_axis(list("abc"))
-        tm.assert_frame_equal(result, expected[0])
-
-        # wrong values for the "axis" parameter
-        for axis in 3, "foo":
-            with pytest.raises(ValueError, match="No axis named"):
-                df.set_axis(list("abc"), axis=axis)

pandas/tests/generic/methods/test_set_axis.py

@@ -0,0 +1,75 @@
+import numpy as np
+import pytest
+
+from pandas import DataFrame, Series
+import pandas._testing as tm
+
+
+class SharedSetAxisTests:
+    @pytest.fixture
+    def obj(self):
+        raise NotImplementedError("Implemented by subclasses")
+
+    def test_set_axis(self, obj):
+        # GH14636; this tests setting index for both Series and DataFrame
+        new_index = list("abcd")[: len(obj)]
+
+        expected = obj.copy()
+        expected.index = new_index
+
+        # inplace=False
+        result = obj.set_axis(new_index, axis=0, inplace=False)
+        tm.assert_equal(expected, result)
+
+    @pytest.mark.parametrize("axis", [0, "index", 1, "columns"])
+    def test_set_axis_inplace_axis(self, axis, obj):
+        # GH#14636
+        if obj.ndim == 1 and axis in [1, "columns"]:
+            # Series only has [0, "index"]
+            return
+
+        new_index = list("abcd")[: len(obj)]
+
+        expected = obj.copy()
+        if axis in [0, "index"]:
+            expected.index = new_index
+        else:
+            expected.columns = new_index
+
+        result = obj.copy()
+        result.set_axis(new_index, axis=axis, inplace=True)
+        tm.assert_equal(result, expected)
+
+    def test_set_axis_unnamed_kwarg_warns(self, obj):
+        # omitting the "axis" parameter
+        new_index = list("abcd")[: len(obj)]
+
+        expected = obj.copy()
+        expected.index = new_index
+
+        with tm.assert_produces_warning(None):
+            result = obj.set_axis(new_index, inplace=False)
+        tm.assert_equal(result, expected)
+
+    @pytest.mark.parametrize("axis", [3, "foo"])
+    def test_set_axis_invalid_axis_name(self, axis, obj):
+        # wrong values for the "axis" parameter
+        with pytest.raises(ValueError, match="No axis named"):
+            obj.set_axis(list("abc"), axis=axis)
+
+
+class TestDataFrameSetAxis(SharedSetAxisTests):
+    @pytest.fixture
+    def obj(self):
+        df = DataFrame(
+            {"A": [1.1, 2.2, 3.3], "B": [5.0, 6.1, 7.2], "C": [4.4, 5.5, 6.6]},
+            index=[2010, 2011, 2012],
+        )
+        return df
+
+
+class TestSeriesSetAxis(SharedSetAxisTests):
+    @pytest.fixture
+    def obj(self):
+        ser = Series(np.arange(4), index=[1, 3, 5, 7], dtype="int64")
+        return ser

pandas/tests/series/test_alter_axes.py

@@ -53,39 +53,3 @@ def test_set_index_makes_timeseries(self):
         s = Series(range(10))
         s.index = idx
         assert s.index.is_all_dates
-
-    def test_set_axis_inplace_axes(self, axis_series):
-        # GH14636
-        ser = Series(np.arange(4), index=[1, 3, 5, 7], dtype="int64")
-
-        expected = ser.copy()
-        expected.index = list("abcd")
-
-        # inplace=True
-        # The FutureWarning comes from the fact that we would like to have
-        # inplace default to False some day
-        result = ser.copy()
-        result.set_axis(list("abcd"), axis=axis_series, inplace=True)
-        tm.assert_series_equal(result, expected)
-
-    def test_set_axis_inplace(self):
-        # GH14636
-
-        s = Series(np.arange(4), index=[1, 3, 5, 7], dtype="int64")
-
-        expected = s.copy()
-        expected.index = list("abcd")
-
-        # inplace=False
-        result = s.set_axis(list("abcd"), axis=0, inplace=False)
-        tm.assert_series_equal(expected, result)
-
-        # omitting the "axis" parameter
-        with tm.assert_produces_warning(None):
-            result = s.set_axis(list("abcd"), inplace=False)
-        tm.assert_series_equal(result, expected)
-
-        # wrong values for the "axis" parameter
-        for axis in [2, "foo"]:
-            with pytest.raises(ValueError, match="No axis named"):
-                s.set_axis(list("abcd"), axis=axis, inplace=False)