Skip to content

Commit 2dcb963

Browse files
authored
TYP: EA.isin (#56423)
1 parent 4f8bb2b commit 2dcb963

File tree

7 files changed

+22
-28
lines changed

7 files changed

+22
-28
lines changed

pandas/core/arrays/arrow/array.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1023,7 +1023,7 @@ def fillna(
10231023

10241024
return super().fillna(value=value, method=method, limit=limit, copy=copy)
10251025

1026-
def isin(self, values) -> npt.NDArray[np.bool_]:
1026+
def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
10271027
# short-circuit to return all False array.
10281028
if not len(values):
10291029
return np.zeros(len(self), dtype=bool)

pandas/core/arrays/base.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1355,15 +1355,15 @@ def equals(self, other: object) -> bool:
13551355
equal_na = self.isna() & other.isna() # type: ignore[operator]
13561356
return bool((equal_values | equal_na).all())
13571357

1358-
def isin(self, values) -> npt.NDArray[np.bool_]:
1358+
def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
13591359
"""
13601360
Pointwise comparison for set containment in the given values.
13611361
13621362
Roughly equivalent to `np.array([x in values for x in self])`
13631363
13641364
Parameters
13651365
----------
1366-
values : Sequence
1366+
values : np.ndarray or ExtensionArray
13671367
13681368
Returns
13691369
-------

pandas/core/arrays/categorical.py

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2570,7 +2570,7 @@ def describe(self) -> DataFrame:
25702570

25712571
return result
25722572

2573-
def isin(self, values) -> npt.NDArray[np.bool_]:
2573+
def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
25742574
"""
25752575
Check whether `values` are contained in Categorical.
25762576
@@ -2580,7 +2580,7 @@ def isin(self, values) -> npt.NDArray[np.bool_]:
25802580
25812581
Parameters
25822582
----------
2583-
values : set or list-like
2583+
values : np.ndarray or ExtensionArray
25842584
The sequence of values to test. Passing in a single string will
25852585
raise a ``TypeError``. Instead, turn a single string into a
25862586
list of one element.
@@ -2611,13 +2611,6 @@ def isin(self, values) -> npt.NDArray[np.bool_]:
26112611
>>> s.isin(['lama'])
26122612
array([ True, False, True, False, True, False])
26132613
"""
2614-
if not is_list_like(values):
2615-
values_type = type(values).__name__
2616-
raise TypeError(
2617-
"only list-like objects are allowed to be passed "
2618-
f"to isin(), you passed a `{values_type}`"
2619-
)
2620-
values = sanitize_array(values, None, None)
26212614
null_mask = np.asarray(isna(values))
26222615
code_values = self.categories.get_indexer_for(values)
26232616
code_values = code_values[null_mask | (code_values >= 0)]

pandas/core/arrays/datetimelike.py

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -734,22 +734,19 @@ def map(self, mapper, na_action=None):
734734
else:
735735
return result.array
736736

737-
def isin(self, values) -> npt.NDArray[np.bool_]:
737+
def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
738738
"""
739739
Compute boolean array of whether each value is found in the
740740
passed set of values.
741741
742742
Parameters
743743
----------
744-
values : set or sequence of values
744+
values : np.ndarray or ExtensionArray
745745
746746
Returns
747747
-------
748748
ndarray[bool]
749749
"""
750-
if not hasattr(values, "dtype"):
751-
values = np.asarray(values)
752-
753750
if values.dtype.kind in "fiuc":
754751
# TODO: de-duplicate with equals, validate_comparison_value
755752
return np.zeros(self.shape, dtype=bool)
@@ -781,15 +778,22 @@ def isin(self, values) -> npt.NDArray[np.bool_]:
781778

782779
if self.dtype.kind in "mM":
783780
self = cast("DatetimeArray | TimedeltaArray", self)
784-
values = values.as_unit(self.unit)
781+
# error: Item "ExtensionArray" of "ExtensionArray | ndarray[Any, Any]"
782+
# has no attribute "as_unit"
783+
values = values.as_unit(self.unit) # type: ignore[union-attr]
785784

786785
try:
787-
self._check_compatible_with(values)
786+
# error: Argument 1 to "_check_compatible_with" of "DatetimeLikeArrayMixin"
787+
# has incompatible type "ExtensionArray | ndarray[Any, Any]"; expected
788+
# "Period | Timestamp | Timedelta | NaTType"
789+
self._check_compatible_with(values) # type: ignore[arg-type]
788790
except (TypeError, ValueError):
789791
# Includes tzawareness mismatch and IncompatibleFrequencyError
790792
return np.zeros(self.shape, dtype=bool)
791793

792-
return isin(self.asi8, values.asi8)
794+
# error: Item "ExtensionArray" of "ExtensionArray | ndarray[Any, Any]"
795+
# has no attribute "asi8"
796+
return isin(self.asi8, values.asi8) # type: ignore[union-attr]
793797

794798
# ------------------------------------------------------------------
795799
# Null Handling

pandas/core/arrays/interval.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1789,12 +1789,8 @@ def contains(self, other):
17891789
other < self._right if self.open_right else other <= self._right
17901790
)
17911791

1792-
def isin(self, values) -> npt.NDArray[np.bool_]:
1793-
if not hasattr(values, "dtype"):
1794-
values = np.array(values)
1795-
values = extract_array(values, extract_numpy=True)
1796-
1797-
if isinstance(values.dtype, IntervalDtype):
1792+
def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
1793+
if isinstance(values, IntervalArray):
17981794
if self.closed != values.closed:
17991795
# not comparable -> no overlap
18001796
return np.zeros(self.shape, dtype=bool)

pandas/core/arrays/masked.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -955,7 +955,7 @@ def take(
955955

956956
# error: Return type "BooleanArray" of "isin" incompatible with return type
957957
# "ndarray" in supertype "ExtensionArray"
958-
def isin(self, values) -> BooleanArray: # type: ignore[override]
958+
def isin(self, values: ArrayLike) -> BooleanArray: # type: ignore[override]
959959
from pandas.core.arrays import BooleanArray
960960

961961
# algorithms.isin will eventually convert values to an ndarray, so no extra

pandas/core/arrays/string_arrow.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
from collections.abc import Sequence
5555

5656
from pandas._typing import (
57+
ArrayLike,
5758
AxisInt,
5859
Dtype,
5960
Scalar,
@@ -212,7 +213,7 @@ def _maybe_convert_setitem_value(self, value):
212213
raise TypeError("Scalar must be NA or str")
213214
return super()._maybe_convert_setitem_value(value)
214215

215-
def isin(self, values) -> npt.NDArray[np.bool_]:
216+
def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
216217
value_set = [
217218
pa_scalar.as_py()
218219
for pa_scalar in [pa.scalar(value, from_pandas=True) for value in values]

0 commit comments

Comments
 (0)