Skip to content

Commit baf1d2c

Browse files
dannyi96dannyi96
authored
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
1 parent 3f9c042 commit baf1d2c

File tree

4 files changed

+91
-5
lines changed

4 files changed

+91
-5
lines changed

‎doc/source/whatsnew/v1.6.0.rst‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ enhancement2
2828

2929
Other enhancements
3030
^^^^^^^^^^^^^^^^^^
31-
-
31+
- :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`)
3232
-
3333

3434
.. ---------------------------------------------------------------------------

‎pandas/core/generic.py‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4619,7 +4619,7 @@ def _update_inplace(self, result, verify_is_copy: bool_t = True) -> None:
46194619
self._maybe_update_cacher(verify_is_copy=verify_is_copy, inplace=True)
46204620

46214621
@final
4622-
def add_prefix(self: NDFrameT, prefix: str) -> NDFrameT:
4622+
def add_prefix(self: NDFrameT, prefix: str, axis: Axis | None = None) -> NDFrameT:
46234623
"""
46244624
Prefix labels with string `prefix`.
46254625
@@ -4630,6 +4630,10 @@ def add_prefix(self: NDFrameT, prefix: str) -> NDFrameT:
46304630
----------
46314631
prefix : str
46324632
The string to add before each label.
4633+
axis : {{0 or 'index', 1 or 'columns', None}}, default None
4634+
Axis to add prefix on
4635+
4636+
.. versionadded:: 1.6.0
46334637
46344638
Returns
46354639
-------
@@ -4675,15 +4679,19 @@ def add_prefix(self: NDFrameT, prefix: str) -> NDFrameT:
46754679
"""
46764680
f = functools.partial("{prefix}{}".format, prefix=prefix)
46774681

4678-
mapper = {self._info_axis_name: f}
4682+
axis_name = self._info_axis_name
4683+
if axis is not None:
4684+
axis_name = self._get_axis_name(axis)
4685+
4686+
mapper = {axis_name: f}
46794687
# error: Incompatible return value type (got "Optional[NDFrameT]",
46804688
# expected "NDFrameT")
46814689
# error: Argument 1 to "rename" of "NDFrame" has incompatible type
46824690
# "**Dict[str, partial[str]]"; expected "Union[str, int, None]"
46834691
return self._rename(**mapper) # type: ignore[return-value, arg-type]
46844692

46854693
@final
4686-
def add_suffix(self: NDFrameT, suffix: str) -> NDFrameT:
4694+
def add_suffix(self: NDFrameT, suffix: str, axis: Axis | None = None) -> NDFrameT:
46874695
"""
46884696
Suffix labels with string `suffix`.
46894697
@@ -4694,6 +4702,10 @@ def add_suffix(self: NDFrameT, suffix: str) -> NDFrameT:
46944702
----------
46954703
suffix : str
46964704
The string to add after each label.
4705+
axis : {{0 or 'index', 1 or 'columns', None}}, default None
4706+
Axis to add suffix on
4707+
4708+
.. versionadded:: 1.6.0
46974709
46984710
Returns
46994711
-------
@@ -4739,7 +4751,11 @@ def add_suffix(self: NDFrameT, suffix: str) -> NDFrameT:
47394751
"""
47404752
f = functools.partial("{}{suffix}".format, suffix=suffix)
47414753

4742-
mapper = {self._info_axis_name: f}
4754+
axis_name = self._info_axis_name
4755+
if axis is not None:
4756+
axis_name = self._get_axis_name(axis)
4757+
4758+
mapper = {axis_name: f}
47434759
# error: Incompatible return value type (got "Optional[NDFrameT]",
47444760
# expected "NDFrameT")
47454761
# error: Argument 1 to "rename" of "NDFrame" has incompatible type

‎pandas/tests/frame/methods/test_add_prefix_suffix.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import pytest
2+
13
from pandas import Index
24
import pandas._testing as tm
35

@@ -18,3 +20,30 @@ def test_add_prefix_suffix(float_frame):
1820
with_pct_suffix = float_frame.add_suffix("%")
1921
expected = Index([f"{c}%" for c in float_frame.columns])
2022
tm.assert_index_equal(with_pct_suffix.columns, expected)
23+
24+
25+
def test_add_prefix_suffix_axis(float_frame):
26+
# GH 47819
27+
with_prefix = float_frame.add_prefix("foo#", axis=0)
28+
expected = Index([f"foo#{c}" for c in float_frame.index])
29+
tm.assert_index_equal(with_prefix.index, expected)
30+
31+
with_prefix = float_frame.add_prefix("foo#", axis=1)
32+
expected = Index([f"foo#{c}" for c in float_frame.columns])
33+
tm.assert_index_equal(with_prefix.columns, expected)
34+
35+
with_pct_suffix = float_frame.add_suffix("#foo", axis=0)
36+
expected = Index([f"{c}#foo" for c in float_frame.index])
37+
tm.assert_index_equal(with_pct_suffix.index, expected)
38+
39+
with_pct_suffix = float_frame.add_suffix("#foo", axis=1)
40+
expected = Index([f"{c}#foo" for c in float_frame.columns])
41+
tm.assert_index_equal(with_pct_suffix.columns, expected)
42+
43+
44+
def test_add_prefix_suffix_invalid_axis(float_frame):
45+
with pytest.raises(ValueError, match="No axis named 2 for object type DataFrame"):
46+
float_frame.add_prefix("foo#", axis=2)
47+
48+
with pytest.raises(ValueError, match="No axis named 2 for object type DataFrame"):
49+
float_frame.add_suffix("foo#", axis=2)

‎pandas/tests/series/methods/test_add_prefix_suffix.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
import pytest
2+
3+
from pandas import Index
4+
import pandas._testing as tm
5+
6+
7+
def test_add_prefix_suffix(string_series):
8+
with_prefix = string_series.add_prefix("foo#")
9+
expected = Index([f"foo#{c}" for c in string_series.index])
10+
tm.assert_index_equal(with_prefix.index, expected)
11+
12+
with_suffix = string_series.add_suffix("#foo")
13+
expected = Index([f"{c}#foo" for c in string_series.index])
14+
tm.assert_index_equal(with_suffix.index, expected)
15+
16+
with_pct_prefix = string_series.add_prefix("%")
17+
expected = Index([f"%{c}" for c in string_series.index])
18+
tm.assert_index_equal(with_pct_prefix.index, expected)
19+
20+
with_pct_suffix = string_series.add_suffix("%")
21+
expected = Index([f"{c}%" for c in string_series.index])
22+
tm.assert_index_equal(with_pct_suffix.index, expected)
23+
24+
25+
def test_add_prefix_suffix_axis(string_series):
26+
# GH 47819
27+
with_prefix = string_series.add_prefix("foo#", axis=0)
28+
expected = Index([f"foo#{c}" for c in string_series.index])
29+
tm.assert_index_equal(with_prefix.index, expected)
30+
31+
with_pct_suffix = string_series.add_suffix("#foo", axis=0)
32+
expected = Index([f"{c}#foo" for c in string_series.index])
33+
tm.assert_index_equal(with_pct_suffix.index, expected)
34+
35+
36+
def test_add_prefix_suffix_invalid_axis(string_series):
37+
with pytest.raises(ValueError, match="No axis named 1 for object type Series"):
38+
string_series.add_prefix("foo#", axis=1)
39+
40+
with pytest.raises(ValueError, match="No axis named 1 for object type Series"):
41+
string_series.add_suffix("foo#", axis=1)

0 commit comments

Comments
 (0)