change how futurewarning is handled in the test

Author: jona-sassenhagen

pandas/core/frame.py

@@ -5526,6 +5526,7 @@ def shift(
                     raise TypeError(
                         f"Periods must be integer, but {period} is {type(period)}."
                     )
+                print(super())
                 shifted_dataframes.append(
                     super()
                     .shift(periods=period, freq=freq, axis=axis, fill_value=fill_value)

pandas/core/groupby/groupby.py

@@ -4866,6 +4866,7 @@ def shift(
         freq=None,
         axis: Axis | lib.NoDefault = lib.no_default,
         fill_value=None,
+        suffix: str | None = None,
     ):
         """
         Shift each group by periods observations.
@@ -4936,17 +4937,38 @@ def shift(
          catfish    NaN  NaN
         goldfish    5.0  8.0
         """
-        if is_list_like(periods):
-            raise NotImplementedError(
-                "shift with multiple periods is not implemented yet for groupby."
-            )
-
         if axis is not lib.no_default:
             axis = self.obj._get_axis_number(axis)
             self._deprecate_axis(axis, "shift")
         else:
             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."
+                )
+            if len(periods) == 0:
+                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}")
+                )
+            return concat(shifted_dataframes, axis=1) if shifted_dataframes else self
+
         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)

pandas/tests/groupby/test_api.py

@@ -192,8 +192,6 @@ def test_frame_consistency(groupby_func):
         exclude_expected = {"numeric_only"}
     elif groupby_func in ("quantile",):
         exclude_expected = {"method", "axis"}
-    elif groupby_func in ("shift",):
-        exclude_expected = {"suffix"}
 
     # Ensure excluded arguments are actually in the signatures
     assert result & exclude_result == exclude_result
@@ -254,8 +252,6 @@ def test_series_consistency(request, groupby_func):
         exclude_expected = {"args", "kwargs"}
     elif groupby_func in ("quantile",):
         exclude_result = {"numeric_only"}
-    elif groupby_func in ("shift",):
-        exclude_expected = {"suffix"}
 
     # Ensure excluded arguments are actually in the signatures
     assert result & exclude_result == exclude_result

pandas/tests/groupby/test_groupby_shift_diff.py

@@ -156,10 +156,16 @@ def test_multindex_empty_shift_with_fill():
     tm.assert_index_equal(shifted.index, shifted_with_fill.index)
 
 
+@pytest.mark.filterwarnings("ignore::FutureWarning")
 def test_group_shift_with_multiple_periods():
     df = DataFrame({"a": [1, 2, 3, 3, 2], "b": [True, True, False, False, True]})
-    with pytest.raises(
-        NotImplementedError,
-        match=r"shift with multiple periods is not implemented yet for groupby.",
-    ):
-        df.groupby("a")["b"].shift([1, 2])
+
+    shifted_df = df.groupby("b")[["a"]].shift([0, 1])
+    expected_df = DataFrame(
+        {"a_0": [1, 2, 3, 3, 2], "a_1": [np.nan, 1.0, np.nan, 3.0, 2.0]}
+    )
+    tm.assert_frame_equal(shifted_df, expected_df)
+
+    # series
+    shifted_series = df.groupby("b")["a"].shift([0, 1])
+    tm.assert_frame_equal(shifted_series, expected_df)