add Hashable type to __get_item__ #592 #596
Conversation
tests/test_frame.py
Outdated
| SECOND = "haydar" | ||
|
|
||
| df = pd.DataFrame(data = [[12.2, 10], [8.8, 15]], columns=[MyEnum.FIRST, MyEnum.SECOND]) | ||
| df[MyEnum.FIRST] |
There was a problem hiding this comment.
Can you use check(assert_type(df[MyEnum.FIRST], pd.Series), pd.Series)
There was a problem hiding this comment.
Thanks I just added it.
pandas-stubs/core/frame.pyi
Outdated
| @@ -516,7 +516,7 @@ class DataFrame(NDFrame, OpsMixin): | |||
| def T(self) -> DataFrame: ... | |||
| def __getattr__(self, name: str) -> Series: ... | |||
| @overload | |||
| def __getitem__(self, idx: Scalar | tuple[Hashable, ...]) -> Series: ... | |||
| def __getitem__(self, idx: Scalar | Hashable | tuple[Hashable, ...]) -> Series: ... | |||
There was a problem hiding this comment.
Could reduce it to def __getitem__(self, idx: Scalar | Hashable) -> Series: ... as a tuple is declared as hashable in typeshed.
(Many members of Scalar are probably also hashable - not sure whether all of them are hashable.)
There was a problem hiding this comment.
Thanks I just removed tuple[Hashable, ...] from the signature.
There was a problem hiding this comment.
Similarly, should we include a list[Hashable] in the 3rd overload?
The 3rd overload:
@overload
def __getitem__(
self,
idx: Series[_bool]
| DataFrame
| Index
| np_ndarray_str
| np_ndarray_bool
| list[_ScalarOrTupleT],
) -> DataFrame: ...There was a problem hiding this comment.
Similarly, should we include a
list[Hashable]in the 3rd overload?
Would need list[HashableT] to simulate the covariant behavior.
Do we have tests for __getitem__ and slice? I believe that should be failing now since slice is hashable. Probably would need to make the slice overload the first overload.
There was a problem hiding this comment.
I'd prefer a more incremental approach here. I'm concerned that we might be making some of the types too wide. Let's have this PR address the issue as reported.
There was a problem hiding this comment.
This is actually more complicated than I thought.
Do we have tests for getitem and slice?
Yes, there is a slice test, but it is without assert types.
pandas-stubs/tests/test_frame.py
Line 156 in fbc5f20
There was a problem hiding this comment.
I just added more assert_types in the test to make sure the other overloads are still working fine.
|
I'm happy to help, thanks @Dr-Irv , @twoertwein for the quick reply and the feedback. |
* add Hashable type to __get_item__ pandas-dev#592 * add assert_type to test_types_getitem_with_hashable * code style: formatting * remove tuple[Hashable, ...] from get_item overload1 * added more assert_type to assure other overloads are ok
assert_type()to assert the type of any return value