GH203 Create new overload for DatetimeIndex

Author: loicdiridollou

pandas-stubs/core/frame.pyi

@@ -1076,17 +1076,29 @@ class DataFrame(NDFrame, OpsMixin):
         dropna: _bool = ...,
     ) -> DataFrameGroupBy[Scalar, Literal[False]]: ...
     @overload
+    def groupby(  # type: ignore[overload-overlap] # pyright: ignore reportOverlappingOverload
+        self,
+        by: DatetimeIndex,
+        axis: AxisIndex | NoDefault = ...,
+        level: IndexLabel | None = ...,
+        as_index: Literal[True] = True,
+        sort: _bool = ...,
+        group_keys: _bool = ...,
+        observed: _bool | NoDefault = ...,
+        dropna: _bool = ...,
+    ) -> DataFrameGroupBy[Timestamp, Literal[True]]: ...
+    @overload
     def groupby(
         self,
         by: DatetimeIndex,
         axis: AxisIndex | NoDefault = ...,
         level: IndexLabel | None = ...,
-        as_index: _bool = ...,
+        as_index: Literal[False] = ...,
         sort: _bool = ...,
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[Timestamp, bool]: ...
+    ) -> DataFrameGroupBy[Timestamp, Literal[False]]: ...
     @overload
     def groupby(
         self,

pandas-stubs/core/groupby/generic.pyi

@@ -30,6 +30,7 @@ from typing_extensions import (
 )
 
 from pandas._libs.lib import NoDefault
+from pandas._libs.tslibs.timestamps import Timestamp
 from pandas._typing import (
     S1,
     AggFuncTypeBase,
@@ -395,3 +396,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
     def size(self: DataFrameGroupBy[ByT, Literal[True]]) -> Series[int]: ...
     @overload
     def size(self: DataFrameGroupBy[ByT, Literal[False]]) -> DataFrame: ...
+    @overload
+    def size(self: DataFrameGroupBy[Timestamp, Literal[True]]) -> Series[int]: ...
+    @overload
+    def size(self: DataFrameGroupBy[Timestamp, Literal[False]]) -> DataFrame: ...

tests/test_frame.py

@@ -1043,6 +1043,27 @@ def test_types_groupby_as_index() -> None:
     )
 
 
+def test_types_groupby_as_index_timestamp() -> None:
+    """Test groupby size with DatetimeIndex."""
+    idx = pd.DatetimeIndex(["2023-10-01", "2023-10-02", "2023-10-01"], name="date")
+    sub_idx = pd.DatetimeIndex(["2023-10-01", "2023-10-02"], name="date")
+    df = pd.DataFrame({"a": [1, 2, 3]}, index=idx)
+    check(
+        assert_type(
+            df.groupby(sub_idx, as_index=False).size(),
+            pd.DataFrame,
+        ),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(
+            df.groupby(sub_idx, as_index=True).size(),
+            "pd.Series[int]",
+        ),
+        pd.Series,
+    )
+
+
 def test_types_groupby_size() -> None:
     """Test for GH886."""
     data = [