Skip to content

Commit 5ab67d7

Browse files
authored
TYP: Update type results for Groupby.count() and Groupby.size() (#45904)
1 parent 4d7a03a commit 5ab67d7

File tree

2 files changed

+8
-6
lines changed

2 files changed

+8
-6
lines changed

pandas/core/groupby/generic.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -240,7 +240,7 @@ def _iterate_slices(self) -> Iterable[Series]:
240240
input="series", examples=_apply_docs["series_examples"]
241241
)
242242
)
243-
def apply(self, func, *args, **kwargs):
243+
def apply(self, func, *args, **kwargs) -> Series:
244244
return super().apply(func, *args, **kwargs)
245245

246246
@doc(_agg_template, examples=_agg_examples_doc, klass="Series")

pandas/core/groupby/groupby.py

+7-5
Original file line numberDiff line numberDiff line change
@@ -924,7 +924,9 @@ def _make_wrapper(self, name: str) -> Callable:
924924
# as are not passed directly but in the grouper
925925
f = getattr(self._obj_with_exclusions, name)
926926
if not isinstance(f, types.MethodType):
927-
return self.apply(lambda self: getattr(self, name))
927+
# error: Incompatible return value type
928+
# (got "NDFrameT", expected "Callable[..., Any]") [return-value]
929+
return cast(Callable, self.apply(lambda self: getattr(self, name)))
928930

929931
f = getattr(type(self._obj_with_exclusions), name)
930932
sig = inspect.signature(f)
@@ -1350,7 +1352,7 @@ def _aggregate_with_numba(self, data, func, *args, engine_kwargs=None, **kwargs)
13501352
input="dataframe", examples=_apply_docs["dataframe_examples"]
13511353
)
13521354
)
1353-
def apply(self, func, *args, **kwargs):
1355+
def apply(self, func, *args, **kwargs) -> NDFrameT:
13541356

13551357
func = com.is_builtin_func(func)
13561358

@@ -1410,7 +1412,7 @@ def _python_apply_general(
14101412
f: Callable,
14111413
data: DataFrame | Series,
14121414
not_indexed_same: bool | None = None,
1413-
) -> DataFrame | Series:
1415+
) -> NDFrameT:
14141416
"""
14151417
Apply function f in python space
14161418
@@ -1812,7 +1814,7 @@ def all(self, skipna: bool = True):
18121814
@final
18131815
@Substitution(name="groupby")
18141816
@Appender(_common_see_also)
1815-
def count(self) -> Series | DataFrame:
1817+
def count(self) -> NDFrameT:
18161818
"""
18171819
Compute count of group, excluding missing values.
18181820
@@ -3467,7 +3469,7 @@ def shift(self, periods=1, freq=None, axis=0, fill_value=None):
34673469
@final
34683470
@Substitution(name="groupby")
34693471
@Appender(_common_see_also)
3470-
def diff(self, periods: int = 1, axis: int = 0) -> Series | DataFrame:
3472+
def diff(self, periods: int = 1, axis: int = 0) -> NDFrameT:
34713473
"""
34723474
First discrete difference of element.
34733475

0 commit comments

Comments
 (0)