Skip to content

Commit 19655c2

Browse files
committed
Add test cases for agg dicts with keys which are sub-types of Hashable and adjust type annotation to support this
1 parent c4e37ab commit 19655c2

File tree

2 files changed

+6
-2
lines changed

2 files changed

+6
-2
lines changed

pandas-stubs/_typing.pyi

+2-2
Original file line numberDiff line numberDiff line change
@@ -127,9 +127,9 @@ F = TypeVar("F", bound=FuncType)
127127
HashableT = TypeVar("HashableT", bound=Hashable)
128128

129129
AggFuncTypeBase: TypeAlias = Union[Callable, str, np.ufunc]
130-
AggFuncTypeDictSeries: TypeAlias = Mapping[Hashable, AggFuncTypeBase]
130+
AggFuncTypeDictSeries: TypeAlias = Mapping[HashableT, AggFuncTypeBase]
131131
AggFuncTypeDictFrame: TypeAlias = Mapping[
132-
Hashable, Union[AggFuncTypeBase, list[AggFuncTypeBase]]
132+
HashableT, Union[AggFuncTypeBase, list[AggFuncTypeBase]]
133133
]
134134
AggFuncTypeSeriesToFrame: TypeAlias = Union[
135135
list[AggFuncTypeBase],

tests/test_frame.py

+4
Original file line numberDiff line numberDiff line change
@@ -670,6 +670,10 @@ def test_types_groupby_agg() -> None:
670670
0: lambda x: x.min(),
671671
}
672672
check(assert_type(df.groupby("col1").agg(agg_dict3), pd.DataFrame), pd.DataFrame)
673+
agg_dict4 = {"col2": "sum"}
674+
check(assert_type(df.groupby("col1").agg(agg_dict4), pd.DataFrame), pd.DataFrame)
675+
agg_dict5 = {0: "sum"}
676+
check(assert_type(df.groupby("col1").agg(agg_dict5), pd.DataFrame), pd.DataFrame)
673677
named_agg = pd.NamedAgg(column="col2", aggfunc="max")
674678
check(
675679
assert_type(df.groupby("col1").agg(new_col=named_agg), pd.DataFrame),

0 commit comments

Comments
 (0)