address comments

Author: jona-sassenhagen

pandas/core/frame.py

@@ -5507,6 +5507,13 @@ def shift(
         fill_value: Hashable = lib.no_default,
         suffix: str | None = None,
     ) -> DataFrame:
+        if freq is not None and fill_value is not lib.no_default:
+            # GH#53832
+            raise ValueError(
+                "Cannot pass both 'freq' and 'fill_value' to "
+                f"{type(self).__name__}.shift"
+            )
+
         axis = self._get_axis_number(axis)
 
         if is_list_like(periods):
@@ -5534,13 +5541,6 @@ def shift(
             return concat(shifted_dataframes, axis=1) if shifted_dataframes else self
         periods = cast(int, periods)
 
-        if freq is not None and fill_value is not lib.no_default:
-            # GH#53832
-            raise ValueError(
-                "Cannot pass both 'freq' and 'fill_value' to "
-                f"{type(self).__name__}.shift"
-            )
-
         ncols = len(self.columns)
         arrays = self._mgr.arrays
         if axis == 1 and periods != 0 and ncols > 0 and freq is None:

pandas/core/groupby/groupby.py

@@ -4949,8 +4949,6 @@ def shift(
             axis = 0
 
         if is_list_like(periods):
-            # periods is not necessarily a list, but otherwise mypy complains.
-            periods = cast(list, periods)
             if axis == 1:
                 raise ValueError(
                     "If `periods` contains multiple shifts, `axis` cannot be 1."
@@ -4959,39 +4957,50 @@ def shift(
                 raise ValueError("If `periods` is an iterable, it cannot be empty.")
             from pandas.core.reshape.concat import concat
 
-            shifted_dataframes = []
-            for period in periods:
-                if not isinstance(period, int):
-                    raise TypeError(
-                        f"Periods must be integer, but {period} is {type(period)}."
-                    )
-                shifted_dataframes.append(
-                    DataFrame(
-                        self.shift(
-                            periods=period, freq=freq, axis=axis, fill_value=fill_value
-                        )
-                    ).add_suffix(f"{suffix}_{period}" if suffix else f"_{period}")
+            add_suffix = True
+        else:
+            periods = [periods]
+            add_suffix = False
+
+        shifted_dataframes = []
+        for period in periods:
+            if not isinstance(period, int):
+                raise TypeError(
+                    f"Periods must be integer, but {period} is {type(period)}."
+                )
+            if freq is not None or axis != 0:
+                f = lambda x: x.shift(period, freq, axis, fill_value)
+                shifted = self._python_apply_general(
+                    f, self._selected_obj, is_transform=True
                 )
-            return concat(shifted_dataframes, axis=1) if shifted_dataframes else self
-        periods = cast(int, periods)
 
-        if freq is not None or axis != 0:
-            f = lambda x: x.shift(periods, freq, axis, fill_value)
-            return self._python_apply_general(f, self._selected_obj, is_transform=True)
+            else:
+                ids, _, ngroups = self.grouper.group_info
+                res_indexer = np.zeros(len(ids), dtype=np.int64)
 
-        ids, _, ngroups = self.grouper.group_info
-        res_indexer = np.zeros(len(ids), dtype=np.int64)
+                libgroupby.group_shift_indexer(res_indexer, ids, ngroups, period)
 
-        libgroupby.group_shift_indexer(res_indexer, ids, ngroups, periods)
+                obj = self._obj_with_exclusions
 
-        obj = self._obj_with_exclusions
+                shifted = obj._reindex_with_indexers(
+                    {self.axis: (obj.axes[self.axis], res_indexer)},
+                    fill_value=fill_value,
+                    allow_dups=True,
+                )
 
-        res = obj._reindex_with_indexers(
-            {self.axis: (obj.axes[self.axis], res_indexer)},
-            fill_value=fill_value,
-            allow_dups=True,
+            if add_suffix:
+                if len(shifted.shape) == 1:
+                    shifted = shifted.to_frame()
+                shifted = shifted.add_suffix(
+                    f"{suffix}_{period}" if suffix else f"_{period}"
+                )
+            shifted_dataframes.append(shifted)
+
+        return (
+            shifted_dataframes[0]
+            if len(shifted_dataframes) == 1
+            else concat(shifted_dataframes, axis=1)
         )
-        return res
 
     @final
     @Substitution(name="groupby")

pandas/tests/frame/methods/test_shift.py

@@ -657,7 +657,7 @@ def test_shift_axis1_many_periods(self):
         shifted2 = df.shift(-6, axis=1, fill_value=None)
         tm.assert_frame_equal(shifted2, expected)
 
-    def test_shift_with_iterable(self):
+    def test_shift_with_iterable_basic_functionality(self):
         # GH#44424
         data = {"a": [1, 2, 3], "b": [4, 5, 6]}
         shifts = [0, 1, 2]
@@ -677,19 +677,47 @@ def test_shift_with_iterable(self):
         )
         tm.assert_frame_equal(expected, shifted)
 
-        # test pd.Series
+    def test_shift_with_iterable_series(self):
+        data = {"a": [1, 2, 3]}
+        shifts = [0, 1, 2]
+
+        df = DataFrame(data)
         s: Series = df["a"]
-        df_one_column: DataFrame = df[["a"]]
-        tm.assert_frame_equal(s.shift(shifts), df_one_column.shift(shifts))
+        tm.assert_frame_equal(s.shift(shifts), df.shift(shifts))
+
+    def test_shift_with_iterable_freq_and_fill_value(self):
+        df = DataFrame(
+            np.random.randn(5), index=date_range("1/1/2000", periods=5, freq="H")
+        )
+
+        tm.assert_frame_equal(
+            # rename because shift with an iterable leads to str column names
+            df.shift([1], fill_value=1).rename(columns=lambda x: int(x[0])),
+            df.shift(1, fill_value=1),
+        )
+
+        tm.assert_frame_equal(
+            df.shift([1], freq="H").rename(columns=lambda x: int(x[0])),
+            df.shift(1, freq="H"),
+        )
+
+        msg = r"Cannot pass both 'freq' and 'fill_value' to.*"
+        with pytest.raises(ValueError, match=msg):
+            df.shift([1, 2], fill_value=1, freq="H")
+
+    def test_shift_with_iterable_check_other_arguments(self):
+        data = {"a": [1, 2], "b": [4, 5]}
+        shifts = [0, 1]
+        df = DataFrame(data)
 
         # test suffix
         columns = df[["a"]].shift(shifts, suffix="_suffix").columns
-        assert columns.tolist() == ["a_suffix_0", "a_suffix_1", "a_suffix_2"]
+        assert columns.tolist() == ["a_suffix_0", "a_suffix_1"]
 
         # check bad inputs when doing multiple shifts
         msg = "If `periods` contains multiple shifts, `axis` cannot be 1."
         with pytest.raises(ValueError, match=msg):
-            df.shift([1, 2], axis=1)
+            df.shift(shifts, axis=1)
 
         msg = "Periods must be integer, but s is <class 'str'>."
         with pytest.raises(TypeError, match=msg):

pandas/tests/groupby/test_groupby_shift_diff.py

@@ -156,7 +156,7 @@ def test_multindex_empty_shift_with_fill():
     tm.assert_index_equal(shifted.index, shifted_with_fill.index)
 
 
-@pytest.mark.filterwarnings("ignore::FutureWarning")
+@pytest.mark.filterwarnings("ignore:The 'axis' keyword in")
 def test_group_shift_with_multiple_periods():
     df = DataFrame({"a": [1, 2, 3, 3, 2], "b": [True, True, False, False, True]})