TYP: annotations in core.groupby

pandas/core/groupby/categorical.py

@@ -1,3 +1,5 @@
+from typing import Optional, Tuple
+
 import numpy as np
 
 from pandas.core.algorithms import unique1d
@@ -6,9 +8,12 @@
     CategoricalDtype,
     recode_for_categories,
 )
+from pandas.core.indexes.api import CategoricalIndex
 
 
-def recode_for_groupby(c: Categorical, sort: bool, observed: bool):
+def recode_for_groupby(
+    c: Categorical, sort: bool, observed: bool
+) -> Tuple[Categorical, Optional[Categorical]]:
     """
     Code the categories to ensure we can groupby for categoricals.
 
@@ -73,7 +78,9 @@ def recode_for_groupby(c: Categorical, sort: bool, observed: bool):
     return c.reorder_categories(cat.categories), None
 
 
-def recode_from_groupby(c: Categorical, sort: bool, ci):
+def recode_from_groupby(
+    c: Categorical, sort: bool, ci: CategoricalIndex
+) -> CategoricalIndex:
     """
     Reverse the codes_to_groupby to account for sort / observed.
 

pandas/core/groupby/generic.py

@@ -83,7 +83,7 @@
 from pandas.plotting import boxplot_frame_groupby
 
 if TYPE_CHECKING:
-    from pandas.core.internals import Block
+    from pandas.core.internals import Block  # noqa:F401
 
 
 NamedAgg = namedtuple("NamedAgg", ["column", "aggfunc"])
@@ -1595,7 +1595,7 @@ def _gotitem(self, key, ndim: int, subset=None):
         Parameters
         ----------
         key : string / list of selections
-        ndim : 1,2
+        ndim : {1, 2}
             requested ndim of result
         subset : object, default None
             subset to act on
@@ -1621,7 +1621,7 @@ def _gotitem(self, key, ndim: int, subset=None):
 
         raise AssertionError("invalid ndim for _gotitem")
 
-    def _wrap_frame_output(self, result, obj) -> DataFrame:
+    def _wrap_frame_output(self, result, obj: DataFrame) -> DataFrame:
         result_index = self.grouper.levels[0]
 
         if self.axis == 0:
@@ -1638,7 +1638,7 @@ def _get_data_to_aggregate(self) -> BlockManager:
         else:
             return obj._mgr
 
-    def _insert_inaxis_grouper_inplace(self, result):
+    def _insert_inaxis_grouper_inplace(self, result: DataFrame):
         # zip in reverse so we can always insert at loc 0
         izip = zip(
             *map(
@@ -1716,7 +1716,7 @@ def _wrap_transformed_output(
 
         return result
 
-    def _wrap_agged_blocks(self, blocks: "Sequence[Block]", items: Index) -> DataFrame:
+    def _wrap_agged_blocks(self, blocks: Sequence["Block"], items: Index) -> DataFrame:
         if not self.as_index:
             index = np.arange(blocks[0].values.shape[-1])
             mgr = BlockManager(blocks, axes=[items, index])
@@ -1743,7 +1743,7 @@ def _iterate_column_groupbys(self):
                 exclusions=self.exclusions,
             )
 
-    def _apply_to_column_groupbys(self, func):
+    def _apply_to_column_groupbys(self, func) -> DataFrame:
         from pandas.core.reshape.concat import concat
 
         return concat(
@@ -1752,7 +1752,7 @@ def _apply_to_column_groupbys(self, func):
             axis=1,
         )
 
-    def count(self):
+    def count(self) -> DataFrame:
         """
         Compute count of group, excluding missing values.
 
@@ -1782,7 +1782,7 @@ def count(self):
 
         return self._reindex_output(result, fill_value=0)
 
-    def nunique(self, dropna: bool = True):
+    def nunique(self, dropna: bool = True) -> DataFrame:
         """
         Return DataFrame with counts of unique elements in each position.
 
@@ -1848,6 +1848,7 @@ def nunique(self, dropna: bool = True):
             ],
             axis=1,
         )
+        assert isinstance(results, DataFrame)  # for mypy
 
         if axis_number == 1:
             results = results.T

pandas/core/groupby/groupby.py

@@ -459,7 +459,7 @@ def f(self):
 
 
 @contextmanager
-def _group_selection_context(groupby):
+def _group_selection_context(groupby: "_GroupBy"):
     """
     Set / reset the _group_selection_context.
     """
@@ -489,7 +489,7 @@ def __init__(
         keys: Optional[_KeysArgType] = None,
         axis: int = 0,
         level=None,
-        grouper: "Optional[ops.BaseGrouper]" = None,
+        grouper: Optional["ops.BaseGrouper"] = None,
         exclusions=None,
         selection=None,
         as_index: bool = True,
@@ -734,7 +734,7 @@ def pipe(self, func, *args, **kwargs):
 
     plot = property(GroupByPlot)
 
-    def _make_wrapper(self, name):
+    def _make_wrapper(self, name: str) -> Callable:
         assert name in self._apply_allowlist
 
         with _group_selection_context(self):

pandas/core/groupby/grouper.py

@@ -568,7 +568,9 @@ def codes(self) -> np.ndarray:
     @cache_readonly
     def result_index(self) -> Index:
         if self.all_grouper is not None:
-            return recode_from_groupby(self.all_grouper, self.sort, self.group_index)
+            group_idx = self.group_index
+            assert isinstance(group_idx, CategoricalIndex)  # set in __init__
+            return recode_from_groupby(self.all_grouper, self.sort, group_idx)
         return self.group_index
 
     @property
@@ -607,7 +609,7 @@ def get_grouper(
     mutated: bool = False,
     validate: bool = True,
     dropna: bool = True,
-) -> "Tuple[ops.BaseGrouper, List[Hashable], FrameOrSeries]":
+) -> Tuple["ops.BaseGrouper", List[Hashable], FrameOrSeries]:
     """
     Create and return a BaseGrouper, which is an internal
     mapping of how to create the grouper indexers.

pandas/core/groupby/ops.py

@@ -82,7 +82,7 @@ class BaseGrouper:
     def __init__(
         self,
         axis: Index,
-        groupings: "Sequence[grouper.Grouping]",
+        groupings: Sequence["grouper.Grouping"],
         sort: bool = True,
         group_keys: bool = True,
         mutated: bool = False,