From ec8c1c4ecf1f0d375d8d1287ee5cbd456852faea Mon Sep 17 00:00:00 2001
From: Richard Shadrach <45562402+rhshadrach@users.noreply.github.com>
Date: Tue, 6 Oct 2020 23:38:49 -0400
Subject: [PATCH] TYP: selection and groups type-hinting in groupby (#36643)

* TYP: selection and groups type-hinting in groupby

* PrettyDict -> Dict

* Cleanup from merge
---
 pandas/core/base.py            |  4 ++--
 pandas/core/groupby/groupby.py | 17 ++++++++++++-----
 pandas/core/groupby/ops.py     |  4 ++--
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/pandas/core/base.py b/pandas/core/base.py
index 9e6f93b656af8..564a0af3527c5 100644
--- a/pandas/core/base.py
+++ b/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)
 
diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
index 887f50f8dbcd5..8340f964fb44b 100644
--- a/pandas/core/groupby/groupby.py
+++ b/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,7 +554,7 @@ 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.
@@ -555,7 +562,7 @@ def _assure_grouper(self):
         pass
 
     @property
-    def groups(self):
+    def groups(self) -> Dict[Hashable, np.ndarray]:
         """
         Dict {group name -> group labels}.
         """
diff --git a/pandas/core/groupby/ops.py b/pandas/core/groupby/ops.py
index 6051aa3022da1..d16ffbca16a3d 100644
--- a/pandas/core/groupby/ops.py
+++ b/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