diff --git a/doc/source/whatsnew/v1.6.0.rst b/doc/source/whatsnew/v1.6.0.rst index 83dfacb46784b..6b7ae767d79d4 100644 --- a/doc/source/whatsnew/v1.6.0.rst +++ b/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`) - .. --------------------------------------------------------------------------- diff --git a/pandas/core/generic.py b/pandas/core/generic.py index abab32ae145bd..f91d15e1a6487 100644 --- a/pandas/core/generic.py +++ b/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,7 +4679,11 @@ 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 @@ -4683,7 +4691,7 @@ def add_prefix(self: NDFrameT, prefix: str) -> NDFrameT: 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 diff --git a/pandas/tests/frame/methods/test_add_prefix_suffix.py b/pandas/tests/frame/methods/test_add_prefix_suffix.py index ea75e9ff51552..92d7cdd7990e1 100644 --- a/pandas/tests/frame/methods/test_add_prefix_suffix.py +++ b/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) diff --git a/pandas/tests/series/methods/test_add_prefix_suffix.py b/pandas/tests/series/methods/test_add_prefix_suffix.py new file mode 100644 index 0000000000000..289a56b98b7e1 --- /dev/null +++ b/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)