Skip to content

Commit ec8c1c4

Browse files
authored
TYP: selection and groups type-hinting in groupby (#36643)
* TYP: selection and groups type-hinting in groupby * PrettyDict -> Dict * Cleanup from merge
1 parent 44ce988 commit ec8c1c4

File tree

3 files changed

+16
-9
lines changed

3 files changed

+16
-9
lines changed

pandas/core/base.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import numpy as np
1010

1111
import pandas._libs.lib as lib
12-
from pandas._typing import AggFuncType, AggFuncTypeBase, Label
12+
from pandas._typing import AggFuncType, AggFuncTypeBase, IndexLabel, Label
1313
from pandas.compat import PYPY
1414
from pandas.compat.numpy import function as nv
1515
from pandas.errors import AbstractMethodError
@@ -136,7 +136,7 @@ class SelectionMixin:
136136
object sub-classes need to define: obj, exclusions
137137
"""
138138

139-
_selection = None
139+
_selection: Optional[IndexLabel] = None
140140
_internal_names = ["_cache", "__setstate__"]
141141
_internal_names_set = set(_internal_names)
142142

pandas/core/groupby/groupby.py

+12-5
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,14 @@ class providing the base-class of operations.
3737

3838
from pandas._libs import Timestamp, lib
3939
import pandas._libs.groupby as libgroupby
40-
from pandas._typing import F, FrameOrSeries, FrameOrSeriesUnion, Label, Scalar
40+
from pandas._typing import (
41+
F,
42+
FrameOrSeries,
43+
FrameOrSeriesUnion,
44+
IndexLabel,
45+
Label,
46+
Scalar,
47+
)
4148
from pandas.compat.numpy import function as nv
4249
from pandas.errors import AbstractMethodError
4350
from pandas.util._decorators import Appender, Substitution, cache_readonly, doc
@@ -487,10 +494,10 @@ def __init__(
487494
obj: FrameOrSeries,
488495
keys: Optional[_KeysArgType] = None,
489496
axis: int = 0,
490-
level=None,
497+
level: Optional[IndexLabel] = None,
491498
grouper: Optional["ops.BaseGrouper"] = None,
492499
exclusions: Optional[Set[Label]] = None,
493-
selection=None,
500+
selection: Optional[IndexLabel] = None,
494501
as_index: bool = True,
495502
sort: bool = True,
496503
group_keys: bool = True,
@@ -547,15 +554,15 @@ def __repr__(self) -> str:
547554
# TODO: Better repr for GroupBy object
548555
return object.__repr__(self)
549556

550-
def _assure_grouper(self):
557+
def _assure_grouper(self) -> None:
551558
"""
552559
We create the grouper on instantiation sub-classes may have a
553560
different policy.
554561
"""
555562
pass
556563

557564
@property
558-
def groups(self):
565+
def groups(self) -> Dict[Hashable, np.ndarray]:
559566
"""
560567
Dict {group name -> group labels}.
561568
"""

pandas/core/groupby/ops.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
"""
88

99
import collections
10-
from typing import List, Optional, Sequence, Tuple, Type
10+
from typing import Dict, Hashable, List, Optional, Sequence, Tuple, Type
1111

1212
import numpy as np
1313

@@ -249,7 +249,7 @@ def size(self) -> Series:
249249
return Series(out, index=self.result_index, dtype="int64")
250250

251251
@cache_readonly
252-
def groups(self):
252+
def groups(self) -> Dict[Hashable, np.ndarray]:
253253
""" dict {group name -> group labels} """
254254
if len(self.groupings) == 1:
255255
return self.groupings[0].groups

0 commit comments

Comments
 (0)