ENH: Add axis parameter to add_prefix and add_suffix (#48085)

* Add axis parameter to add_prefix and add_suffix

* Add axis parameter to add_prefix and add_suffix

* Added testcases

* docstring update

* docstring update

* updated whatsnew file

* updated whatsnew file

* updated whatsnew file

* review comments

* address merge conflixts

* address merge conflixts

* updated whatsnew

* review comments

* review comments

* review comments

Author: dannyi96

doc/source/whatsnew/v1.6.0.rst

@@ -28,7 +28,7 @@ enhancement2
 
 Other enhancements
 ^^^^^^^^^^^^^^^^^^
--
+- :meth:`Series.add_suffix`, :meth:`DataFrame.add_suffix`, :meth:`Series.add_prefix` and :meth:`DataFrame.add_prefix` support an ``axis`` argument. If ``axis`` is set, the default behaviour of which axis to consider can be overwritten (:issue:`47819`)
 -
 
 .. ---------------------------------------------------------------------------

pandas/core/generic.py

@@ -4619,7 +4619,7 @@ def _update_inplace(self, result, verify_is_copy: bool_t = True) -> None:
         self._maybe_update_cacher(verify_is_copy=verify_is_copy, inplace=True)
 
     @final
-    def add_prefix(self: NDFrameT, prefix: str) -> NDFrameT:
+    def add_prefix(self: NDFrameT, prefix: str, axis: Axis | None = None) -> NDFrameT:
         """
         Prefix labels with string `prefix`.
 
@@ -4630,6 +4630,10 @@ def add_prefix(self: NDFrameT, prefix: str) -> NDFrameT:
         ----------
         prefix : str
             The string to add before each label.
+        axis : {{0 or 'index', 1 or 'columns', None}}, default None
+            Axis to add prefix on
+
+             .. versionadded:: 1.6.0
 
         Returns
         -------
@@ -4675,15 +4679,19 @@ def add_prefix(self: NDFrameT, prefix: str) -> NDFrameT:
         """
         f = functools.partial("{prefix}{}".format, prefix=prefix)
 
-        mapper = {self._info_axis_name: f}
+        axis_name = self._info_axis_name
+        if axis is not None:
+            axis_name = self._get_axis_name(axis)
+
+        mapper = {axis_name: f}
         # error: Incompatible return value type (got "Optional[NDFrameT]",
         # expected "NDFrameT")
         # error: Argument 1 to "rename" of "NDFrame" has incompatible type
         # "**Dict[str, partial[str]]"; expected "Union[str, int, None]"
         return self._rename(**mapper)  # type: ignore[return-value, arg-type]
 
     @final
-    def add_suffix(self: NDFrameT, suffix: str) -> NDFrameT:
+    def add_suffix(self: NDFrameT, suffix: str, axis: Axis | None = None) -> NDFrameT:
         """
         Suffix labels with string `suffix`.
 
@@ -4694,6 +4702,10 @@ def add_suffix(self: NDFrameT, suffix: str) -> NDFrameT:
         ----------
         suffix : str
             The string to add after each label.
+        axis : {{0 or 'index', 1 or 'columns', None}}, default None
+            Axis to add suffix on
+
+             .. versionadded:: 1.6.0
 
         Returns
         -------
@@ -4739,7 +4751,11 @@ def add_suffix(self: NDFrameT, suffix: str) -> NDFrameT:
         """
         f = functools.partial("{}{suffix}".format, suffix=suffix)
 
-        mapper = {self._info_axis_name: f}
+        axis_name = self._info_axis_name
+        if axis is not None:
+            axis_name = self._get_axis_name(axis)
+
+        mapper = {axis_name: f}
         # error: Incompatible return value type (got "Optional[NDFrameT]",
         # expected "NDFrameT")
         # error: Argument 1 to "rename" of "NDFrame" has incompatible type

pandas/tests/frame/methods/test_add_prefix_suffix.py

@@ -1,3 +1,5 @@
+import pytest
+
 from pandas import Index
 import pandas._testing as tm
 
@@ -18,3 +20,30 @@ def test_add_prefix_suffix(float_frame):
     with_pct_suffix = float_frame.add_suffix("%")
     expected = Index([f"{c}%" for c in float_frame.columns])
     tm.assert_index_equal(with_pct_suffix.columns, expected)
+
+
+def test_add_prefix_suffix_axis(float_frame):
+    # GH 47819
+    with_prefix = float_frame.add_prefix("foo#", axis=0)
+    expected = Index([f"foo#{c}" for c in float_frame.index])
+    tm.assert_index_equal(with_prefix.index, expected)
+
+    with_prefix = float_frame.add_prefix("foo#", axis=1)
+    expected = Index([f"foo#{c}" for c in float_frame.columns])
+    tm.assert_index_equal(with_prefix.columns, expected)
+
+    with_pct_suffix = float_frame.add_suffix("#foo", axis=0)
+    expected = Index([f"{c}#foo" for c in float_frame.index])
+    tm.assert_index_equal(with_pct_suffix.index, expected)
+
+    with_pct_suffix = float_frame.add_suffix("#foo", axis=1)
+    expected = Index([f"{c}#foo" for c in float_frame.columns])
+    tm.assert_index_equal(with_pct_suffix.columns, expected)
+
+
+def test_add_prefix_suffix_invalid_axis(float_frame):
+    with pytest.raises(ValueError, match="No axis named 2 for object type DataFrame"):
+        float_frame.add_prefix("foo#", axis=2)
+
+    with pytest.raises(ValueError, match="No axis named 2 for object type DataFrame"):
+        float_frame.add_suffix("foo#", axis=2)

pandas/tests/series/methods/test_add_prefix_suffix.py

@@ -0,0 +1,41 @@
+import pytest
+
+from pandas import Index
+import pandas._testing as tm
+
+
+def test_add_prefix_suffix(string_series):
+    with_prefix = string_series.add_prefix("foo#")
+    expected = Index([f"foo#{c}" for c in string_series.index])
+    tm.assert_index_equal(with_prefix.index, expected)
+
+    with_suffix = string_series.add_suffix("#foo")
+    expected = Index([f"{c}#foo" for c in string_series.index])
+    tm.assert_index_equal(with_suffix.index, expected)
+
+    with_pct_prefix = string_series.add_prefix("%")
+    expected = Index([f"%{c}" for c in string_series.index])
+    tm.assert_index_equal(with_pct_prefix.index, expected)
+
+    with_pct_suffix = string_series.add_suffix("%")
+    expected = Index([f"{c}%" for c in string_series.index])
+    tm.assert_index_equal(with_pct_suffix.index, expected)
+
+
+def test_add_prefix_suffix_axis(string_series):
+    # GH 47819
+    with_prefix = string_series.add_prefix("foo#", axis=0)
+    expected = Index([f"foo#{c}" for c in string_series.index])
+    tm.assert_index_equal(with_prefix.index, expected)
+
+    with_pct_suffix = string_series.add_suffix("#foo", axis=0)
+    expected = Index([f"{c}#foo" for c in string_series.index])
+    tm.assert_index_equal(with_pct_suffix.index, expected)
+
+
+def test_add_prefix_suffix_invalid_axis(string_series):
+    with pytest.raises(ValueError, match="No axis named 1 for object type Series"):
+        string_series.add_prefix("foo#", axis=1)
+
+    with pytest.raises(ValueError, match="No axis named 1 for object type Series"):
+        string_series.add_suffix("foo#", axis=1)