TYP: selection and groups type-hinting in groupby (pandas-dev#36643)

* TYP: selection and groups type-hinting in groupby

* PrettyDict -> Dict

* Cleanup from merge

pandas/core/base.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 import pandas._libs.lib as lib
-from pandas._typing import AggFuncType, AggFuncTypeBase, Label
+from pandas._typing import AggFuncType, AggFuncTypeBase, IndexLabel, Label
 from pandas.compat import PYPY
 from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
@@ -136,7 +136,7 @@ class SelectionMixin:
     object sub-classes need to define: obj, exclusions
     """
 
-    _selection = None
+    _selection: Optional[IndexLabel] = None
     _internal_names = ["_cache", "__setstate__"]
     _internal_names_set = set(_internal_names)
 

pandas/core/groupby/groupby.py

@@ -37,7 +37,14 @@ class providing the base-class of operations.
 
 from pandas._libs import Timestamp, lib
 import pandas._libs.groupby as libgroupby
-from pandas._typing import F, FrameOrSeries, FrameOrSeriesUnion, Label, Scalar
+from pandas._typing import (
+    F,
+    FrameOrSeries,
+    FrameOrSeriesUnion,
+    IndexLabel,
+    Label,
+    Scalar,
+)
 from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import Appender, Substitution, cache_readonly, doc
@@ -487,10 +494,10 @@ def __init__(
         obj: FrameOrSeries,
         keys: Optional[_KeysArgType] = None,
         axis: int = 0,
-        level=None,
+        level: Optional[IndexLabel] = None,
         grouper: Optional["ops.BaseGrouper"] = None,
         exclusions: Optional[Set[Label]] = None,
-        selection=None,
+        selection: Optional[IndexLabel] = None,
         as_index: bool = True,
         sort: bool = True,
         group_keys: bool = True,
@@ -547,15 +554,15 @@ def __repr__(self) -> str:
         # TODO: Better repr for GroupBy object
         return object.__repr__(self)
 
-    def _assure_grouper(self):
+    def _assure_grouper(self) -> None:
         """
         We create the grouper on instantiation sub-classes may have a
         different policy.
         """
         pass
 
     @property
-    def groups(self):
+    def groups(self) -> Dict[Hashable, np.ndarray]:
         """
         Dict {group name -> group labels}.
         """

pandas/core/groupby/ops.py

@@ -7,7 +7,7 @@
 """
 
 import collections
-from typing import List, Optional, Sequence, Tuple, Type
+from typing import Dict, Hashable, List, Optional, Sequence, Tuple, Type
 
 import numpy as np
 
@@ -249,7 +249,7 @@ def size(self) -> Series:
         return Series(out, index=self.result_index, dtype="int64")
 
     @cache_readonly
-    def groups(self):
+    def groups(self) -> Dict[Hashable, np.ndarray]:
         """ dict {group name -> group labels} """
         if len(self.groupings) == 1:
             return self.groupings[0].groups