mostly types, also exclude groupby

Author: jona-sassenhagen

pandas/core/dtypes/dtypes.py

@@ -951,7 +951,7 @@ class PeriodDtype(PeriodDtypeBase, PandasExtensionDtype):
     # "Dict[int, PandasExtensionDtype]", base class "PandasExtensionDtype"
     # defined the type as "Dict[str, PandasExtensionDtype]")  [assignment]
     _cache_dtypes: dict[BaseOffset, PeriodDtype] = {}  # type: ignore[assignment] # noqa: E501
-    __hash__ = PeriodDtypeBase.__hash__
+    __hash__ = PeriodDtypeBase.__hash__  # type: ignore[assignment]
     _freq: BaseOffset
 
     def __new__(cls, freq):

pandas/core/frame.py

@@ -5510,6 +5510,8 @@ def shift(
         axis = self._get_axis_number(axis)
 
         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."
@@ -5518,18 +5520,20 @@ def shift(
                 raise ValueError("If `periods` is an iterable, it cannot be empty.")
             from pandas.core.reshape.concat import concat
 
-            result = []
+            shifted_dataframes = []
             for period in periods:
-                if not isinstance(period, int):
+                if not isinstance(int, period):
                     raise TypeError(
                         f"Periods must be integer, but {period} is {type(period)}."
                     )
-                result.append(
+                period = cast(int, period)
+                shifted_dataframes.append(
                     super()
                     .shift(periods=period, freq=freq, axis=axis, fill_value=fill_value)
                     .add_suffix(f"{suffix}_{period}" if suffix else f"_{period}")
                 )
-            return concat(result, axis=1) if result else self
+            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

pandas/core/generic.py

@@ -10522,7 +10522,7 @@ def mask(
     @doc(klass=_shared_doc_kwargs["klass"])
     def shift(
         self,
-        periods: int | Iterable = 1,
+        periods: int | Iterable[int] = 1,
         freq=None,
         axis: Axis = 0,
         fill_value: Hashable = lib.no_default,
@@ -10656,6 +10656,7 @@ def shift(
             return self.to_frame().shift(
                 periods=periods, freq=freq, axis=axis, fill_value=fill_value
             )
+        periods = cast(int, periods)
 
         if freq is None:
             # when freq is None, data is shifted, index is not

pandas/core/groupby/groupby.py

@@ -4866,7 +4866,6 @@ def shift(
         freq=None,
         axis: Axis | lib.NoDefault = lib.no_default,
         fill_value=None,
-        suffix: str | None = None,
     ):
         """
         Shift each group by periods observations.
@@ -4875,8 +4874,9 @@ def shift(
 
         Parameters
         ----------
-        periods : int, default 1
-            Number of periods to shift.
+        periods : int | Iterable[int], default 1
+            Number of periods to shift. If a list of values, shift each group by
+            each period.
         freq : str, optional
             Frequency string.
         axis : axis to shift, default 0
@@ -4888,8 +4888,6 @@ def shift(
 
         fill_value : optional
             The scalar value to use for newly introduced missing values.
-        suffix : str, optional
-            An optional suffix to append when there are multiple periods.
 
         Returns
         -------
@@ -4938,6 +4936,11 @@ 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")

pandas/tests/groupby/test_api.py

@@ -192,6 +192,8 @@ 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
@@ -252,6 +254,8 @@ 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

@@ -154,3 +154,12 @@ def test_multindex_empty_shift_with_fill():
     shifted_with_fill = df.groupby(["a", "b"]).shift(1, fill_value=0)
     tm.assert_frame_equal(shifted, shifted_with_fill)
     tm.assert_index_equal(shifted.index, shifted_with_fill.index)
+
+
+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])