Skip to content

Commit 849d4ac

Browse files
author
Mateusz
committed
TYP: Fix core/groupby/generic.py type ignores
GH37715
1 parent 6d0e480 commit 849d4ac

File tree

1 file changed

+16
-24
lines changed

1 file changed

+16
-24
lines changed

pandas/core/groupby/generic.py

+16-24
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
Mapping,
1919
NamedTuple,
2020
Sequence,
21+
TYPE_CHECKING,
2122
TypeVar,
2223
Union,
2324
cast,
@@ -88,6 +89,9 @@
8889

8990
from pandas.plotting import boxplot_frame_groupby
9091

92+
if TYPE_CHECKING:
93+
from pandas._libs import Interval
94+
9195
# TODO(typing) the return value on this callable should be any *scalar*.
9296
AggScalar = Union[str, Callable[..., Any]]
9397
# TODO: validate types on ScalarResult and move to _typing
@@ -270,9 +274,7 @@ def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs)
270274
func = maybe_mangle_lambdas(func)
271275
ret = self._aggregate_multiple_funcs(func)
272276
if relabeling:
273-
# error: Incompatible types in assignment (expression has type
274-
# "Optional[List[str]]", variable has type "Index")
275-
ret.columns = columns # type: ignore[assignment]
277+
ret.columns = Index(columns)
276278
return ret
277279

278280
else:
@@ -622,33 +624,27 @@ def value_counts(
622624
ids, val = ids[mask], val[mask]
623625

624626
if bins is None:
625-
lab, lev = algorithms.factorize(val, sort=True)
627+
lab, lev = cast(
628+
tuple[np.ndarray, Index], algorithms.factorize(val, sort=True)
629+
)
626630
llab = lambda lab, inc: lab[inc]
627631
else:
628632

629-
# lab is a Categorical with categories an IntervalIndex
630-
lab = cut(Series(val), bins, include_lowest=True)
631-
# error: "ndarray" has no attribute "cat"
632-
lev = lab.cat.categories # type: ignore[attr-defined]
633-
# error: No overload variant of "take" of "_ArrayOrScalarCommon" matches
634-
# argument types "Any", "bool", "Union[Any, float]"
635-
lab = lev.take( # type: ignore[call-overload]
636-
# error: "ndarray" has no attribute "cat"
637-
lab.cat.codes, # type: ignore[attr-defined]
633+
# labels is a Series with categories an IntervalIndex
634+
labels: Series = cut(Series(val), bins, include_lowest=True)
635+
lev = labels.cat.categories
636+
lab = lev.take(
637+
labels.cat.codes,
638638
allow_fill=True,
639-
# error: Item "ndarray" of "Union[ndarray, Index]" has no attribute
640-
# "_na_value"
641-
fill_value=lev._na_value, # type: ignore[union-attr]
639+
fill_value=lev._na_value,
642640
)
643641
llab = lambda lab, inc: lab[inc]._multiindex.codes[-1]
644642

645643
if is_interval_dtype(lab.dtype):
646644
# TODO: should we do this inside II?
647645

648-
# error: "ndarray" has no attribute "left"
649-
# error: "ndarray" has no attribute "right"
650646
sorter = np.lexsort(
651-
(lab.left, lab.right, ids) # type: ignore[attr-defined]
647+
(cast(Interval, lab).left, cast(Interval, lab).right, ids)
652648
)
653649
else:
654650
sorter = np.lexsort((lab, ids))
@@ -675,11 +671,7 @@ def value_counts(
675671
# multi-index components
676672
codes = self.grouper.reconstructed_codes
677673
codes = [rep(level_codes) for level_codes in codes] + [llab(lab, inc)]
678-
# error: List item 0 has incompatible type "Union[ndarray[Any, Any], Index]";
679-
# expected "Index"
680-
levels = [ping.group_index for ping in self.grouper.groupings] + [
681-
lev # type: ignore[list-item]
682-
]
674+
levels = [ping.group_index for ping in self.grouper.groupings] + [lev]
683675

684676
if dropna:
685677
mask = codes[-1] != -1

0 commit comments

Comments
 (0)