TYP: Annotate groupby/ops.py

pandas/core/groupby/ops.py

@@ -7,14 +7,14 @@
 """
 
 import collections
-from typing import List, Optional, Sequence, Tuple, Type
+from typing import Any, Callable, List, Optional, Sequence, Tuple, Type
 
 import numpy as np
 
 from pandas._libs import NaT, iNaT, lib
 import pandas._libs.groupby as libgroupby
 import pandas._libs.reduction as libreduction
-from pandas._typing import FrameOrSeries
+from pandas._typing import FrameOrSeries, Label
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import cache_readonly
 
@@ -102,7 +102,7 @@ def groupings(self) -> List["grouper.Grouping"]:
         return self._groupings
 
     @property
-    def shape(self):
+    def shape(self) -> Tuple[int, ...]:
         return tuple(ping.ngroups for ping in self.groupings)
 
     def __iter__(self):
@@ -148,7 +148,9 @@ def _get_group_keys(self):
             # provide "flattened" iterator for multi-group setting
             return get_flattened_iterator(comp_ids, ngroups, self.levels, self.codes)
 
-    def apply(self, f, data: FrameOrSeries, axis: int = 0):
+    def apply(
+        self, f: Callable[[FrameOrSeries], Any], data: FrameOrSeries, axis: int = 0
+    ):
         mutated = self.mutated
         splitter = self._get_splitter(data, axis=axis)
         group_keys = self._get_group_keys()
@@ -229,7 +231,7 @@ def levels(self) -> List[Index]:
         return [ping.group_index for ping in self.groupings]
 
     @property
-    def names(self):
+    def names(self) -> List[Label]:
         return [ping.name for ping in self.groupings]
 
     def size(self) -> Series:
@@ -307,7 +309,7 @@ def result_index(self) -> Index:
         )
         return result
 
-    def get_group_levels(self):
+    def get_group_levels(self) -> List[Index]:
         if not self.compressed and len(self.groupings) == 1:
             return [self.groupings[0].result_index]
 
@@ -356,7 +358,9 @@ def _is_builtin_func(self, arg):
         """
         return SelectionMixin._builtin_table.get(arg, arg)
 
-    def _get_cython_function(self, kind: str, how: str, values, is_numeric: bool):
+    def _get_cython_function(
+        self, kind: str, how: str, values: np.ndarray, is_numeric: bool
+    ):
 
         dtype_str = values.dtype.name
         ftype = self._cython_functions[kind][how]
@@ -425,7 +429,7 @@ def _get_cython_func_and_vals(
         return func, values
 
     def _cython_operation(
-        self, kind: str, values, how: str, axis, min_count: int = -1, **kwargs
+        self, kind: str, values, how: str, axis: int, min_count: int = -1, **kwargs,
     ) -> Tuple[np.ndarray, Optional[List[str]]]:
         """
         Returns the values of a cython operation as a Tuple of [data, names].
@@ -578,13 +582,13 @@ def _cython_operation(
         return result, names
 
     def aggregate(
-        self, values, how: str, axis: int = 0, min_count: int = -1
+        self, values: np.ndarray, how: str, axis: int = 0, min_count: int = -1
     ) -> Tuple[np.ndarray, Optional[List[str]]]:
         return self._cython_operation(
             "aggregate", values, how, axis, min_count=min_count
         )
 
-    def transform(self, values, how: str, axis: int = 0, **kwargs):
+    def transform(self, values: np.ndarray, how: str, axis: int = 0, **kwargs):
         return self._cython_operation("transform", values, how, axis, **kwargs)
 
     def _aggregate(
@@ -608,7 +612,7 @@ def _transform(
 
         return result
 
-    def agg_series(self, obj: Series, func):
+    def agg_series(self, obj: Series, func: Callable[[Series], Any]):
         # Caller is responsible for checking ngroups != 0
         assert self.ngroups != 0
 
@@ -637,7 +641,7 @@ def agg_series(self, obj: Series, func):
                 raise
         return self._aggregate_series_pure_python(obj, func)
 
-    def _aggregate_series_fast(self, obj: Series, func):
+    def _aggregate_series_fast(self, obj: Series, func: Callable[[Series], Any]):
         # At this point we have already checked that
         #  - obj.index is not a MultiIndex
         #  - obj is backed by an ndarray, not ExtensionArray
@@ -656,7 +660,7 @@ def _aggregate_series_fast(self, obj: Series, func):
         result, counts = grouper.get_result()
         return result, counts
 
-    def _aggregate_series_pure_python(self, obj: Series, func):
+    def _aggregate_series_pure_python(self, obj: Series, func: Callable[[Series], Any]):
 
         group_index, _, ngroups = self.group_info
 
@@ -828,11 +832,11 @@ def result_index(self):
         return self.binlabels
 
     @property
-    def levels(self):
+    def levels(self) -> List:
         return [self.binlabels]
 
     @property
-    def names(self):
+    def names(self) -> List:
         return [self.binlabels.name]
 
     @property
@@ -842,7 +846,7 @@ def groupings(self) -> "List[grouper.Grouping]":
             for lvl, name in zip(self.levels, self.names)
         ]
 
-    def agg_series(self, obj: Series, func):
+    def agg_series(self, obj: Series, func: Callable[[Series], Any]):
         # Caller is responsible for checking ngroups != 0
         assert self.ngroups != 0
         assert len(self.bins) > 0  # otherwise we'd get IndexError in get_result
@@ -916,7 +920,9 @@ def _chop(self, sdata: Series, slice_obj: slice) -> Series:
 
 
 class FrameSplitter(DataSplitter):
-    def fast_apply(self, f, sdata: FrameOrSeries, names):
+    def fast_apply(
+        self, f: Callable[[FrameOrSeries], Any], sdata: FrameOrSeries, names
+    ):
         # must return keys::list, values::list, mutated::bool
         starts, ends = lib.generate_slices(self.slabels, self.ngroups)
         return libreduction.apply_frame_axis0(sdata, f, names, starts, ends)