TYP: Fix core/groupby/generic.py type ignores

GH37715

Author: Mateusz

pandas/core/groupby/generic.py

@@ -11,13 +11,15 @@
 from functools import partial
 from textwrap import dedent
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Hashable,
     Iterable,
     Mapping,
     NamedTuple,
     Sequence,
+    Tuple,
     TypeVar,
     Union,
     cast,
@@ -88,6 +90,9 @@
 
 from pandas.plotting import boxplot_frame_groupby
 
+if TYPE_CHECKING:
+    from pandas._libs import Interval
+
 # TODO(typing) the return value on this callable should be any *scalar*.
 AggScalar = Union[str, Callable[..., Any]]
 # TODO: validate types on ScalarResult and move to _typing
@@ -270,9 +275,7 @@ def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs)
             func = maybe_mangle_lambdas(func)
             ret = self._aggregate_multiple_funcs(func)
             if relabeling:
-                # error: Incompatible types in assignment (expression has type
-                # "Optional[List[str]]", variable has type "Index")
-                ret.columns = columns  # type: ignore[assignment]
+                ret.columns = Index(columns)
             return ret
 
         else:
@@ -622,33 +625,27 @@ def value_counts(
         ids, val = ids[mask], val[mask]
 
         if bins is None:
-            lab, lev = algorithms.factorize(val, sort=True)
+            lab, lev = cast(
+                Tuple[np.ndarray, Index], algorithms.factorize(val, sort=True)
+            )
             llab = lambda lab, inc: lab[inc]
         else:
 
-            # lab is a Categorical with categories an IntervalIndex
-            lab = cut(Series(val), bins, include_lowest=True)
-            # error: "ndarray" has no attribute "cat"
-            lev = lab.cat.categories  # type: ignore[attr-defined]
-            # error: No overload variant of "take" of "_ArrayOrScalarCommon" matches
-            # argument types "Any", "bool", "Union[Any, float]"
-            lab = lev.take(  # type: ignore[call-overload]
-                # error: "ndarray" has no attribute "cat"
-                lab.cat.codes,  # type: ignore[attr-defined]
+            # labels is a Series with categories an IntervalIndex
+            labels: Series = cut(Series(val), bins, include_lowest=True)
+            lev = labels.cat.categories
+            lab = lev.take(
+                labels.cat.codes,
                 allow_fill=True,
-                # error: Item "ndarray" of "Union[ndarray, Index]" has no attribute
-                # "_na_value"
-                fill_value=lev._na_value,  # type: ignore[union-attr]
+                fill_value=lev._na_value,
             )
             llab = lambda lab, inc: lab[inc]._multiindex.codes[-1]
 
         if is_interval_dtype(lab.dtype):
             # TODO: should we do this inside II?
 
-            # error: "ndarray" has no attribute "left"
-            # error: "ndarray" has no attribute "right"
             sorter = np.lexsort(
-                (lab.left, lab.right, ids)  # type: ignore[attr-defined]
+                (cast(Interval, lab).left, cast(Interval, lab).right, ids)
             )
         else:
             sorter = np.lexsort((lab, ids))
@@ -675,11 +672,7 @@ def value_counts(
         # multi-index components
         codes = self.grouper.reconstructed_codes
         codes = [rep(level_codes) for level_codes in codes] + [llab(lab, inc)]
-        # error: List item 0 has incompatible type "Union[ndarray[Any, Any], Index]";
-        # expected "Index"
-        levels = [ping.group_index for ping in self.grouper.groupings] + [
-            lev  # type: ignore[list-item]
-        ]
+        levels = [ping.group_index for ping in self.grouper.groupings] + [lev]
 
         if dropna:
             mask = codes[-1] != -1