diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 6cd3236bb..f413e6110 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1009,7 +1009,7 @@ class DataFrame(NDFrame, OpsMixin): self, by: Scalar, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -1022,7 +1022,7 @@ class DataFrame(NDFrame, OpsMixin): self, by: DatetimeIndex, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -1035,7 +1035,7 @@ class DataFrame(NDFrame, OpsMixin): self, by: TimedeltaIndex, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -1048,7 +1048,7 @@ class DataFrame(NDFrame, OpsMixin): self, by: PeriodIndex, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -1061,7 +1061,7 @@ class DataFrame(NDFrame, OpsMixin): self, by: IntervalIndex[IntervalT], axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -1074,7 +1074,7 @@ class DataFrame(NDFrame, OpsMixin): self, by: MultiIndex | GroupByObjectNonScalar | None = ..., axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -1087,7 +1087,7 @@ class DataFrame(NDFrame, OpsMixin): self, by: Series[SeriesByT], axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -1100,7 +1100,7 @@ class DataFrame(NDFrame, OpsMixin): self, by: CategoricalIndex | Index | Series, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index aeee0fecd..2a42e8a3d 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -117,6 +117,7 @@ from pandas._typing import ( HashableT3, IgnoreRaise, IndexingInt, + IndexLabel, IntDtypeArg, InterpolateOptions, IntervalClosedType, @@ -547,7 +548,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self, by: Scalar, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -560,7 +561,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self, by: DatetimeIndex, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -573,7 +574,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self, by: TimedeltaIndex, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -586,7 +587,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self, by: PeriodIndex, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -599,7 +600,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self, by: IntervalIndex[IntervalT], axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -612,7 +613,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self, by: MultiIndex | GroupByObjectNonScalar = ..., axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -625,7 +626,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self, by: Series[SeriesByT], axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., @@ -638,7 +639,7 @@ class Series(IndexOpsMixin[S1], NDFrame): self, by: CategoricalIndex | Index | Series, axis: AxisIndex = ..., - level: Level | None = ..., + level: IndexLabel | None = ..., as_index: _bool = ..., sort: _bool = ..., group_keys: _bool = ..., diff --git a/tests/test_frame.py b/tests/test_frame.py index a25e9b683..bdd9bf491 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -1073,6 +1073,21 @@ def test_types_groupby_iter() -> None: ) +def test_types_groupby_level() -> None: + # GH 836 + data = { + "col1": [0, 0, 0], + "col2": [0, 1, 0], + "col3": [1, 2, 3], + "col4": [1, 2, 3], + } + df = pd.DataFrame(data=data).set_index(["col1", "col2", "col3"]) + check( + assert_type(df.groupby(level=["col1", "col2"]).sum(), pd.DataFrame), + pd.DataFrame, + ) + + def test_types_merge() -> None: df = pd.DataFrame(data={"col1": [1, 1, 2], "col2": [3, 4, 5]}) df2 = pd.DataFrame(data={"col1": [1, 1, 2], "col2": [0, 1, 0]}) diff --git a/tests/test_series.py b/tests/test_series.py index 72ecd0cbe..ba6ae57f7 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -439,6 +439,19 @@ def test_types_max() -> None: s.max(skipna=False) +def test_types_groupby_level() -> None: + # GH 836 + index = pd.MultiIndex.from_tuples( + [(0, 0, 1), (0, 1, 2), (0, 0, 3)], names=["col1", "col2", "col3"] + ) + s = pd.Series([1, 2, 3], index=index) + check( + assert_type(s.groupby(level=["col1", "col2"]).sum(), "pd.Series[int]"), + pd.Series, + np.integer, + ) + + def test_types_quantile() -> None: s = pd.Series([1, 2, 3, 10]) s.quantile([0.25, 0.5])