type isin, remove redundant Iterable | Series (#1507)

* type `isin`, remove redundant `Iterable | Series`

* Iterable -> Iterable[Any], extra test cases

* add overloads for multiindex.isin

* fix pyrefly

* fix pyright

* fix pyrefly

* extra test case, use Collection

* remove extra test case

* extra test cases

* mypy

* ty

Author: MarcoGorelli

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(

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,

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],

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,

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:

tests/series/test_series.py

@@ -1606,6 +1606,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])]

tests/test_frame.py

@@ -2,6 +2,7 @@
 
 from collections import (
     OrderedDict,
+    UserDict,
     UserList,
     defaultdict,
     deque,
@@ -2936,6 +2937,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)