diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 344121c01..d4bb41511 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -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( diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 2c7d9aeb0..09bb37a8a 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -576,7 +576,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, diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index 64c0c85c9..dbb9b766e 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -1,5 +1,6 @@ from collections.abc import ( Callable, + Collection, Hashable, Iterable, Mapping, @@ -165,7 +166,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( # ty: ignore[invalid-method-override] # pyright: ignore[reportIncompatibleMethodOverride] + self, values: Collection[Iterable[Any]], level: None = None + ) -> np_1darray_bool: ... def set_names( self, names: Hashable | Sequence[Hashable] | Mapping[Any, Hashable], diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index dfa25dc93..4f6ddc425 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1305,7 +1305,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, diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 55abdae9a..eda08354e 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -57,8 +57,18 @@ 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) + check(assert_type(mi.isin({iter([3])}), np_1darray_bool), np_1darray_bool) + if TYPE_CHECKING_INVALID_USAGE: + mi.isin({3}) # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + mi.isin(iter([[3]])) # type: ignore[call-overload] # pyright: ignore[reportArgumentType] def test_index_astype() -> None: diff --git a/tests/series/test_series.py b/tests/series/test_series.py index 004fddbb5..941697a95 100644 --- a/tests/series/test_series.py +++ b/tests/series/test_series.py @@ -1607,6 +1607,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])] diff --git a/tests/test_frame.py b/tests/test_frame.py index 32529509f..4780d2af1 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2,6 +2,7 @@ from collections import ( OrderedDict, + UserDict, UserList, defaultdict, deque, @@ -2942,6 +2943,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)