diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 8053c17437c5e..319cda773b693 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -5866,7 +5866,9 @@ def shift( freq: Frequency | None = None, axis: Axis = 0, fill_value: Hashable = lib.no_default, + prefix: str | None = None, suffix: str | None = None, + sep: str ="_", ) -> DataFrame: if freq is not None and fill_value is not lib.no_default: # GH#53832 @@ -5899,11 +5901,17 @@ def shift( shifted_dataframes.append( super() .shift(periods=period, freq=freq, axis=axis, fill_value=fill_value) + .add_prefix(f"{prefix}{sep}" if prefix else "") .add_suffix(f"{suffix}_{period}" if suffix else f"_{period}") + .rename( + columns=lambda col: col.replace(f"{sep}{sep}", sep) + ) ) return concat(shifted_dataframes, axis=1) elif suffix: raise ValueError("Cannot specify `suffix` if `periods` is an int.") + elif prefix: + raise ValueError("Cannot specify `prefix` if `periods` is an int.") periods = cast(int, periods) ncols = len(self.columns) diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 8aae4609b1833..bf3ff1f954741 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -10112,13 +10112,17 @@ def mask( ) @doc(klass=_shared_doc_kwargs["klass"]) + if self.ndim == 1 and (prefix or suffix): + raise ValueError("`prefix` and `suffix` are only supported on DataFrame.shift when `periods` is a list. ") def shift( self, periods: int | Sequence[int] = 1, freq=None, axis: Axis = 0, fill_value: Hashable = lib.no_default, + prefix: str | None = None, suffix: str | None = None, + sep: str = "_" ) -> Self | DataFrame: """ Shift index by desired number of periods with an optional time `freq`. diff --git a/pandas/tests/frame/methods/test_shift.py b/pandas/tests/frame/methods/test_shift.py index b52240c208493..00de4653813bf 100644 --- a/pandas/tests/frame/methods/test_shift.py +++ b/pandas/tests/frame/methods/test_shift.py @@ -725,6 +725,16 @@ def test_shift_with_iterable_check_other_arguments(self): expected = DataFrame({"a_suffix_0": [1, 2], "a_suffix_1": [np.nan, 1.0]}) tm.assert_frame_equal(shifted, expected) + # test prefix and sep + shifted = df[["a"]].shift(shifts, prefix="pre", suffix="suf", sep="-") + expected = DataFrame( + { + "pre-a-suf-0": [1, 2], + "pre-a-suf-1": [np.nan, 1.0] + } + ) + tm.assert_frame_equal(shifted, expected) + # check bad inputs when doing multiple shifts msg = "If `periods` contains multiple shifts, `axis` cannot be 1." with pytest.raises(ValueError, match=msg): diff --git a/pandas/tests/series/methods/test_diff.py b/pandas/tests/series/methods/test_diff.py index 89c34aeb64206..1b5d9d2de1bb0 100644 --- a/pandas/tests/series/methods/test_diff.py +++ b/pandas/tests/series/methods/test_diff.py @@ -1,5 +1,6 @@ import numpy as np import pytest +import pandas as pd from pandas import ( Series, @@ -89,3 +90,12 @@ def test_diff_object_dtype(self): result = ser.diff() expected = ser - ser.shift(1) tm.assert_series_equal(result, expected) + +def test_series_shift_rejects_prefix_suffix(): + s = pd.Series([1, 2, 3]) + + with pytest.raises(ValueError, match="prefix.*DataFrame.shift"): + s.shift(1, prefix="PRE") + + with pytest.raises(ValueError, match="suffix.*DataFrame.shift"): + s.shift(1, suffix="POST")