diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index df25e8b90..d03189dca 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -225,27 +225,6 @@ class _iLocIndexerFrame(_iLocIndexer, Generic[_T]): ) -> None: ... class _LocIndexerFrame(_LocIndexer, Generic[_T]): - @overload - def __getitem__(self, idx: Scalar) -> Series | _T: ... - @overload - def __getitem__( # type: ignore[overload-overlap] - self, - idx: ( - IndexType - | MaskType - | Callable[[DataFrame], IndexType | MaskType | Sequence[Hashable]] - | list[HashableT] - | tuple[ - IndexType - | MaskType - | list[HashableT] - | slice - | _IndexSliceTuple - | Callable, - MaskType | list[HashableT] | IndexType | Callable, - ] - ), - ) -> _T: ... @overload def __getitem__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] self, @@ -277,7 +256,28 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]): ), ) -> Series: ... @overload + def __getitem__(self, idx: Scalar) -> Series | _T: ... + @overload def __getitem__(self, idx: tuple[Scalar, slice]) -> Series | _T: ... + @overload + def __getitem__( + self, + idx: ( + IndexType + | MaskType + | Callable[[DataFrame], IndexType | MaskType | Sequence[Hashable]] + | list[HashableT] + | tuple[ + IndexType + | MaskType + | list[HashableT] + | slice + | _IndexSliceTuple + | Callable, + MaskType | Iterable[HashableT] | IndexType | Callable, + ] + ), + ) -> _T: ... # Keep in sync with `DataFrame.__setitem__` @overload diff --git a/tests/test_frame.py b/tests/test_frame.py index 0a8a33511..e4b8797a3 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2,7 +2,9 @@ from collections import ( OrderedDict, + UserList, defaultdict, + deque, ) from collections.abc import ( Callable, @@ -3718,12 +3720,28 @@ def test_loc_int_set() -> None: df.loc[np.uint64(1)] = [2, 3] -def test_loclist() -> None: - # GH 189 +@pytest.mark.parametrize("col", [1, None]) +@pytest.mark.parametrize("typ", [list, tuple, deque, UserList, iter]) +def test_loc_iterable(col: Hashable, typ: type) -> None: + # GH 189, GH 1410 df = pd.DataFrame({1: [1, 2], None: 5}, columns=pd.Index([1, None], dtype=object)) + check(df.loc[:, typ([col])], pd.DataFrame) - check(assert_type(df.loc[:, [None]], pd.DataFrame), pd.DataFrame) - check(assert_type(df.loc[:, [1]], pd.DataFrame), pd.DataFrame) + if TYPE_CHECKING: + assert_type(df.loc[:, [None]], pd.DataFrame) + assert_type(df.loc[:, [1]], pd.DataFrame) + + assert_type(df.loc[:, (None,)], pd.DataFrame) + assert_type(df.loc[:, (1,)], pd.DataFrame) + + assert_type(df.loc[:, deque([None])], pd.DataFrame) + assert_type(df.loc[:, deque([1])], pd.DataFrame) + + assert_type(df.loc[:, UserList([None])], pd.DataFrame) + assert_type(df.loc[:, UserList([1])], pd.DataFrame) + + assert_type(df.loc[:, (None for _ in [0])], pd.DataFrame) + assert_type(df.loc[:, (1 for _ in [0])], pd.DataFrame) def test_dict_items() -> None: