From ae0174009b694568c288f30be0458b3e225beb76 Mon Sep 17 00:00:00 2001 From: Loic Diridollou Date: Wed, 16 Oct 2024 20:15:57 -0400 Subject: [PATCH] Experiment for size --- pandas-stubs/core/frame.pyi | 31 +++++++++------------------ pandas-stubs/core/groupby/generic.pyi | 11 ++++++++-- pandas-stubs/core/groupby/groupby.pyi | 12 +---------- 3 files changed, 20 insertions(+), 34 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 3d6e222f..e846bbbb 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -12,9 +12,7 @@ from re import Pattern from typing import ( Any, ClassVar, - Generic, Literal, - TypeVar, overload, ) @@ -27,10 +25,7 @@ from pandas import ( ) from pandas.core.arraylike import OpsMixin from pandas.core.generic import NDFrame -from pandas.core.groupby.generic import ( - DataFrameGroupBy, - SeriesGroupBy, -) +from pandas.core.groupby.generic import DataFrameGroupBy from pandas.core.groupby.grouper import Grouper from pandas.core.indexers import BaseIndexer from pandas.core.indexes.base import Index @@ -79,7 +74,6 @@ from pandas._typing import ( Axis, AxisColumn, AxisIndex, - ByT, CalculationMethod, ColspaceArgType, CompressionOptions, @@ -235,11 +229,6 @@ class _LocIndexerFrame(_LocIndexer): value: Scalar | NAType | NaTType | ArrayLike | Series | list | None, ) -> None: ... -_TT = TypeVar("TT", bound=Literal[True, False]) - -class DataFrameGroupByGen(DataFrameGroupBy[ByT], Generic[ByT, _TT]): ... -class SeriesGroupByGen(SeriesGroupBy, Generic[_TT, ByT]): ... - class DataFrame(NDFrame, OpsMixin): __hash__: ClassVar[None] # type: ignore[assignment] @@ -1073,7 +1062,7 @@ class DataFrame(NDFrame, OpsMixin): group_keys: _bool = ..., observed: _bool | NoDefault = ..., dropna: _bool = ..., - ) -> DataFrameGroupByGen[Scalar, Literal[True]]: ... + ) -> DataFrameGroupBy[Scalar, Literal[True]]: ... @overload def groupby( self, @@ -1085,7 +1074,7 @@ class DataFrame(NDFrame, OpsMixin): group_keys: _bool = ..., observed: _bool | NoDefault = ..., dropna: _bool = ..., - ) -> DataFrameGroupByGen[Scalar, Literal[False]]: ... + ) -> DataFrameGroupBy[Scalar, Literal[False]]: ... @overload def groupby( self, @@ -1097,7 +1086,7 @@ class DataFrame(NDFrame, OpsMixin): group_keys: _bool = ..., observed: _bool | NoDefault = ..., dropna: _bool = ..., - ) -> DataFrameGroupBy[Timestamp]: ... + ) -> DataFrameGroupBy[Timestamp, bool]: ... @overload def groupby( self, @@ -1109,7 +1098,7 @@ class DataFrame(NDFrame, OpsMixin): group_keys: _bool = ..., observed: _bool | NoDefault = ..., dropna: _bool = ..., - ) -> DataFrameGroupBy[Timedelta]: ... + ) -> DataFrameGroupBy[Timedelta, bool]: ... @overload def groupby( self, @@ -1121,7 +1110,7 @@ class DataFrame(NDFrame, OpsMixin): group_keys: _bool = ..., observed: _bool | NoDefault = ..., dropna: _bool = ..., - ) -> DataFrameGroupBy[Period]: ... + ) -> DataFrameGroupBy[Period, bool]: ... @overload def groupby( self, @@ -1133,7 +1122,7 @@ class DataFrame(NDFrame, OpsMixin): group_keys: _bool = ..., observed: _bool | NoDefault = ..., dropna: _bool = ..., - ) -> DataFrameGroupBy[IntervalT]: ... + ) -> DataFrameGroupBy[IntervalT, bool]: ... @overload def groupby( self, @@ -1145,7 +1134,7 @@ class DataFrame(NDFrame, OpsMixin): group_keys: _bool = ..., observed: _bool | NoDefault = ..., dropna: _bool = ..., - ) -> DataFrameGroupBy[tuple]: ... + ) -> DataFrameGroupBy[tuple, bool]: ... @overload def groupby( self, @@ -1157,7 +1146,7 @@ class DataFrame(NDFrame, OpsMixin): group_keys: _bool = ..., observed: _bool | NoDefault = ..., dropna: _bool = ..., - ) -> DataFrameGroupBy[SeriesByT]: ... + ) -> DataFrameGroupBy[SeriesByT, bool]: ... @overload def groupby( self, @@ -1169,7 +1158,7 @@ class DataFrame(NDFrame, OpsMixin): group_keys: _bool = ..., observed: _bool | NoDefault = ..., dropna: _bool = ..., - ) -> DataFrameGroupBy[Any]: ... + ) -> DataFrameGroupBy[Any, bool]: ... def pivot( self, *, diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index 5ccc4179..56f8fdb8 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/pandas-stubs/core/groupby/generic.pyi @@ -11,6 +11,7 @@ from typing import ( Generic, Literal, NamedTuple, + TypeVar, final, overload, ) @@ -182,7 +183,9 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]): self, ) -> Iterator[tuple[ByT, Series[S1]]]: ... -class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT]): +_TT = TypeVar("_TT", bound=Literal[True, False]) + +class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]): # error: Overload 3 for "apply" will never be used because its parameters overlap overload 1 @overload # type: ignore[override] def apply( # type: ignore[overload-overlap] @@ -236,7 +239,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT]): @overload def __getitem__( # pyright: ignore[reportIncompatibleMethodOverride] self, key: Iterable[Hashable] | slice - ) -> DataFrameGroupBy[ByT]: ... + ) -> DataFrameGroupBy[ByT, bool]: ... def nunique(self, dropna: bool = ...) -> DataFrame: ... def idxmax( self, @@ -388,3 +391,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT]): def __iter__( # pyright: ignore[reportIncompatibleMethodOverride] self, ) -> Iterator[tuple[ByT, DataFrame]]: ... + @overload + def size(self: DataFrameGroupBy[ByT, Literal[True]]) -> Series[int]: ... + @overload + def size(self: DataFrameGroupBy[ByT, Literal[False]]) -> DataFrame: ... diff --git a/pandas-stubs/core/groupby/groupby.pyi b/pandas-stubs/core/groupby/groupby.pyi index beb03d6d..c638dc28 100644 --- a/pandas-stubs/core/groupby/groupby.pyi +++ b/pandas-stubs/core/groupby/groupby.pyi @@ -18,10 +18,7 @@ from typing import ( import numpy as np from pandas.core.base import SelectionMixin -from pandas.core.frame import ( - DataFrame, - DataFrameGroupByGen, -) +from pandas.core.frame import DataFrame from pandas.core.groupby import ( generic, ops, @@ -56,7 +53,6 @@ from pandas._typing import ( AnyArrayLike, Axis, AxisInt, - ByT, CalculationMethod, Dtype, Frequency, @@ -236,13 +232,7 @@ class GroupBy(BaseGroupBy[NDFrameT]): def sem( self: GroupBy[DataFrame], ddof: int = ..., numeric_only: bool = ... ) -> DataFrame: ... - @final - @overload def size(self: GroupBy[Series]) -> Series[int]: ... - @overload - def size(self: DataFrameGroupByGen[ByT, Literal[True]]) -> Series[int]: ... # type: ignore[misc] - @overload - def size(self: DataFrameGroupByGen[ByT, Literal[False]]) -> DataFrame: ... # type: ignore[misc] @final def sum( self,