Skip to content

Commit c76c147

Browse files
committed
fix issues with Index.__getitem__
1 parent 63d03bc commit c76c147

File tree

3 files changed

+22
-3
lines changed

3 files changed

+22
-3
lines changed

tests/pandas/test_indexes.py

+19
Original file line numberDiff line numberDiff line change
@@ -40,3 +40,22 @@ def test_index_tolist() -> None:
4040
i1 = pd.Index([1, 2, 3])
4141
l1 = i1.tolist()
4242
i2 = i1.to_list()
43+
44+
45+
def test_column_getitem() -> None:
46+
# https://github.com/microsoft/python-type-stubs/issues/199#issuecomment-1132806594
47+
df = pd.DataFrame([[1, 2, 3]], columns=["a", "b", "c"])
48+
49+
column = df.columns[0]
50+
a = df[column]
51+
52+
def test_column_contains() -> None:
53+
# https://github.com/microsoft/python-type-stubs/issues/199
54+
df = pd.DataFrame({'A': [1, 2], 'B': ['c', 'd'], 'E': [3, 4]})
55+
56+
collist = [column for column in df.columns]
57+
58+
collist2 = [column for column in df.columns[df.columns.str.contains('A|B')]]
59+
60+
length = len(df.columns[df.columns.str.contains('A|B')])
61+

typings/pandas/core/frame.pyi

+1-1
Original file line numberDiff line numberDiff line change
@@ -356,7 +356,7 @@ class DataFrame(NDFrame, OpsMixin):
356356
Series[_bool],
357357
DataFrame,
358358
List[_str],
359-
List[Hashable],
359+
List[Scalar],
360360
Index,
361361
np_ndarray_str,
362362
np_ndarray_bool,

typings/pandas/core/indexes/base.pyi

+2-2
Original file line numberDiff line numberDiff line change
@@ -174,11 +174,11 @@ class Index(IndexOpsMixin, PandasObject):
174174
def __hash__(self) -> int: ...
175175
def __setitem__(self, key, value) -> None: ...
176176
@overload
177-
def __getitem__(self, idx: Union[slice, np_ndarray_int64, Index]) -> Index: ...
177+
def __getitem__(self, idx: Union[slice, np_ndarray_int64, Index, Series[bool]]) -> Index: ...
178178
@overload
179179
def __getitem__(
180180
self, idx: Union[int, Tuple[np_ndarray_int64, ...]]
181-
) -> Hashable: ...
181+
) -> Scalar: ...
182182
def append(self, other): ...
183183
def putmask(self, mask, value): ...
184184
def equals(self, other) -> bool: ...

0 commit comments

Comments
 (0)