Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -1783,7 +1783,9 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
axis: Axis = 0,
copy: _bool = True,
) -> Self: ...
def isin(self, values: Iterable | Series | DataFrame | dict) -> Self: ...
def isin(
self, values: Iterable[Any] | Mapping[Hashable, Iterable[Any]] | DataFrame
) -> Self: ...
@property
def plot(self) -> PlotAccessor: ...
def hist(
Expand Down
4 changes: 3 additions & 1 deletion pandas-stubs/core/indexes/base.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -575,7 +575,9 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
def map(
self, mapper: Renamer, na_action: Literal["ignore"] | None = None
) -> Index: ...
def isin(self, values, level=...) -> np_1darray_bool: ...
def isin(
self, values: Iterable[Any], level: Level | None = None
) -> np_1darray_bool: ...
def slice_indexer(
self,
start: Label | None = None,
Expand Down
9 changes: 8 additions & 1 deletion pandas-stubs/core/indexes/multi.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,14 @@ class MultiIndex(Index):
def equal_levels(self, other): ...
def insert(self, loc, item): ...
def delete(self, loc): ...
def isin(self, values, level=...) -> np_1darray_bool: ...
@overload # type: ignore[override]
def isin( # pyrefly: ignore[bad-override]
self, values: Iterable[Any], level: Level
) -> np_1darray_bool: ...
@overload
def isin( # pyright: ignore[reportIncompatibleMethodOverride]
self, values: Iterable[Iterable[Any]], level: None = None
) -> np_1darray_bool: ...
def set_names(
self,
names: Hashable | Sequence[Hashable] | Mapping[Any, Hashable],
Expand Down
2 changes: 1 addition & 1 deletion pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -1302,7 +1302,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
show_counts: bool | None = ...,
) -> None: ...
def memory_usage(self, index: _bool = True, deep: _bool = False) -> int: ...
def isin(self, values: Iterable | Series[S1] | dict) -> Series[_bool]: ...
def isin(self, values: Iterable[Any]) -> Series[_bool]: ...
def between(
self,
left: Scalar | ListLikeU,
Expand Down
12 changes: 10 additions & 2 deletions tests/indexes/test_indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,16 @@ def test_index_duplicated() -> None:

def test_index_isin() -> None:
ind = pd.Index([1, 2, 3, 4, 5])
isin = ind.isin([2, 4])
check(assert_type(isin, np_1darray_bool), np_1darray_bool)
check(assert_type(ind.isin([2, 4]), np_1darray_bool), np_1darray_bool)
check(assert_type(ind.isin({2, 4}), np_1darray_bool), np_1darray_bool)
check(assert_type(ind.isin(pd.Series([2, 4])), np_1darray_bool), np_1darray_bool)
check(assert_type(ind.isin(pd.Index([2, 4])), np_1darray_bool), np_1darray_bool)
check(assert_type(ind.isin(iter([2, "4"])), np_1darray_bool), np_1darray_bool)

mi = pd.MultiIndex.from_arrays([[1, 2, 3]])
check(assert_type(mi.isin([[3]]), np_1darray_bool), np_1darray_bool)
if TYPE_CHECKING_INVALID_USAGE:
mi.isin({3}) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]


def test_index_astype() -> None:
Expand Down
11 changes: 11 additions & 0 deletions tests/series/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -1610,6 +1610,17 @@ def test_series_min_max_sub_axis() -> None:
check(assert_type(df.max(axis=1), pd.Series), pd.Series)


def test_series_isin() -> None:
s = pd.Series([1, 2, 3, 4, 5])
check(assert_type(s.isin([3, 4]), "pd.Series[bool]"), pd.Series, np.bool_)
check(assert_type(s.isin({3, 4}), "pd.Series[bool]"), pd.Series, np.bool_)
check(
assert_type(s.isin(pd.Series([3, 4])), "pd.Series[bool]"), pd.Series, np.bool_
)
check(assert_type(s.isin(pd.Index([3, 4])), "pd.Series[bool]"), pd.Series, np.bool_)
check(assert_type(s.isin(iter([3, "4"])), "pd.Series[bool]"), pd.Series, np.bool_)


def test_series_index_isin() -> None:
s = pd.Series([1, 2, 3, 4, 5], index=[1, 2, 2, 3, 3])
t1 = s.loc[s.index.isin([1, 3])]
Expand Down
15 changes: 15 additions & 0 deletions tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from collections import (
OrderedDict,
UserDict,
UserList,
defaultdict,
deque,
Expand Down Expand Up @@ -2940,6 +2941,20 @@ def test_getmultiindex_columns() -> None:
check(assert_type(df[li[0]], pd.Series), pd.Series)


def test_frame_isin() -> None:
df = pd.DataFrame({"x": [1, 2, 3, 4, 5]}, index=[1, 2, 3, 4, 5])
check(assert_type(df.isin([1, 3, 5]), pd.DataFrame), pd.DataFrame)
check(assert_type(df.isin({1, 3, 5}), pd.DataFrame), pd.DataFrame)
check(assert_type(df.isin(pd.Series([1, 3, 5])), pd.DataFrame), pd.DataFrame)
check(assert_type(df.isin(pd.Index([1, 3, 5])), pd.DataFrame), pd.DataFrame)
check(assert_type(df.isin(df), pd.DataFrame), pd.DataFrame)
check(assert_type(df.isin({"x": [1, 2]}), pd.DataFrame), pd.DataFrame)
check(
assert_type(df.isin(UserDict({"x": iter([1, "2"])})), pd.DataFrame),
pd.DataFrame,
)


def test_frame_getitem_isin() -> None:
df = pd.DataFrame({"x": [1, 2, 3, 4, 5]}, index=[1, 2, 3, 4, 5])
check(assert_type(df[df.index.isin([1, 3, 5])], pd.DataFrame), pd.DataFrame)
Expand Down