Extends DataFrame.groupby overloads to recognize some scalar index types

pandas-stubs/core/frame.pyi

@@ -27,11 +27,15 @@ from pandas.core.arraylike import OpsMixin
 from pandas.core.generic import NDFrame
 from pandas.core.groupby.generic import (
     _DataFrameGroupByNonScalar,
+    _DataFrameGroupByPeriod,
     _DataFrameGroupByScalar,
 )
 from pandas.core.groupby.grouper import Grouper
 from pandas.core.indexers import BaseIndexer
 from pandas.core.indexes.base import Index
+from pandas.core.indexes.datetimes import DatetimeIndex
+from pandas.core.indexes.period import PeriodIndex
+from pandas.core.indexes.timedeltas import TimedeltaIndex
 from pandas.core.indexing import (
     _iLocIndexer,
     _IndexSliceTuple,
@@ -52,6 +56,7 @@ import xarray as xr
 
 from pandas._libs.missing import NAType
 from pandas._libs.tslibs import BaseOffset
+from pandas._libs.tslibs.period import Period
 from pandas._typing import (
     S1,
     AggFuncTypeBase,
@@ -1000,9 +1005,9 @@ class DataFrame(NDFrame, OpsMixin):
         errors: IgnoreRaise = ...,
     ) -> None: ...
     @overload
-    def groupby(
+    def groupby(  # type: ignore[misc]
         self,
-        by: Scalar,
+        by: Scalar | DatetimeIndex | TimedeltaIndex,
         axis: Axis = ...,
         level: Level | None = ...,
         as_index: _bool = ...,
@@ -1013,6 +1018,19 @@ class DataFrame(NDFrame, OpsMixin):
         dropna: _bool = ...,
     ) -> _DataFrameGroupByScalar: ...
     @overload
+    def groupby(  # type: ignore[misc]  # pyright: ignore[reportOverlappingOverload]
+        self,
+        by: PeriodIndex,
+        axis: Axis = ...,
+        level: Level | None = ...,
+        as_index: _bool = ...,
+        sort: _bool = ...,
+        group_keys: _bool = ...,
+        squeeze: _bool = ...,
+        observed: _bool = ...,
+        dropna: _bool = ...,
+    ) -> _DataFrameGroupByPeriod: ...
+    @overload
     def groupby(
         self,
         by: GroupByObjectNonScalar | None = ...,

pandas-stubs/core/groupby/generic.pyi

@@ -25,6 +25,7 @@ from pandas.core.groupby.grouper import Grouper
 from pandas.core.series import Series
 from typing_extensions import TypeAlias
 
+from pandas._libs.tslibs.period import Period
 from pandas._typing import (
     S1,
     AggFuncTypeBase,
@@ -149,6 +150,9 @@ class SeriesGroupBy(GroupBy, Generic[S1]):
 class _DataFrameGroupByScalar(DataFrameGroupBy):
     def __iter__(self) -> Iterator[tuple[Scalar, DataFrame]]: ...
 
+class _DataFrameGroupByPeriod(DataFrameGroupBy):
+    def __iter__(self) -> Iterator[tuple[Period, DataFrame]]: ...
+
 class _DataFrameGroupByNonScalar(DataFrameGroupBy):
     def __iter__(self) -> Iterator[tuple[tuple, DataFrame]]: ...
 

tests/test_frame.py

@@ -1980,6 +1980,46 @@ def test_groupby_result() -> None:
         pass
 
 
+def test_groupby_result_for_scalar_indeces() -> None:
+    # GH 674
+    df = pd.DataFrame({"date": pd.date_range("2020-01-01", "2020-12-31"), "days": 1})
+    period_index = pd.PeriodIndex(df.date, freq="M")
+    iterator = df.groupby(period_index).__iter__()
+    assert_type(iterator, Iterator[Tuple[pd.Period, pd.DataFrame]])
+    index, value = next(iterator)
+    assert_type((index, value), Tuple[pd.Period, pd.DataFrame])
+
+    check(assert_type(index, pd.Period), pd.Period)
+    check(assert_type(value, pd.DataFrame), pd.DataFrame)
+
+    dt_index = pd.DatetimeIndex(df.date)
+    iterator2 = df.groupby(dt_index).__iter__()
+    assert_type(iterator2, Iterator[Tuple[Scalar, pd.DataFrame]])
+    index2, value2 = next(iterator2)
+    assert_type((index2, value2), Tuple[Scalar, pd.DataFrame])
+
+    check(assert_type(index2, Scalar), pd.Timestamp)
+    check(assert_type(value2, pd.DataFrame), pd.DataFrame)
+
+    tdelta_index = pd.TimedeltaIndex(df.date - pd.Timestamp("2020-01-01"))
+    iterator3 = df.groupby(tdelta_index).__iter__()
+    assert_type(iterator3, Iterator[Tuple[Scalar, pd.DataFrame]])
+    index3, value3 = next(iterator3)
+    assert_type((index3, value3), Tuple[Scalar, pd.DataFrame])
+
+    check(assert_type(index3, Scalar), pd.Timedelta)
+    check(assert_type(value3, pd.DataFrame), pd.DataFrame)
+
+    for p, g in df.groupby(period_index):
+        pass
+
+    for dt, g in df.groupby(dt_index):
+        pass
+
+    for tdelta, g in df.groupby(tdelta_index):
+        pass
+
+
 def test_setitem_list():
     # GH 153
     lst1: list[str] = ["a", "b", "c"]