TYP: ExtensionArray delete() and searchsorted()

pandas/_typing.py

@@ -91,6 +91,20 @@
 PandasScalar = Union["Period", "Timestamp", "Timedelta", "Interval"]
 Scalar = Union[PythonScalar, PandasScalar]
 
+# numpy compatible types
+NumpyValueArrayLike = Union[PythonScalar, Sequence[PythonScalar]]
+NumpySorter = Union[
+    int,
+    np.integer,
+    bool,
+    np.ndarray,
+    Sequence[int],
+    Sequence[np.integer],
+    Sequence[bool],
+    None,
+]
+
+
 # timestamp and timedelta convertible types
 
 TimestampConvertibleTypes = Union[

pandas/core/algorithms.py

@@ -30,6 +30,8 @@
     ArrayLike,
     DtypeObj,
     FrameOrSeriesUnion,
+    NumpySorter,
+    NumpyValueArrayLike,
     Scalar,
 )
 from pandas.util._decorators import doc
@@ -1507,7 +1509,12 @@ def take(
 # ------------ #
 
 
-def searchsorted(arr, value, side="left", sorter=None) -> np.ndarray:
+def searchsorted(
+    arr: ArrayLike,
+    value: NumpyValueArrayLike,
+    side: Literal["left", "right"] = "left",
+    sorter: NumpySorter = None,
+) -> np.ndarray:
     """
     Find indices where elements should be inserted to maintain order.
 
@@ -1532,7 +1539,7 @@ def searchsorted(arr, value, side="left", sorter=None) -> np.ndarray:
         Input array. If `sorter` is None, then it must be sorted in
         ascending order, otherwise `sorter` must be an array of indices
         that sort it.
-    value : array_like
+    value : array_like or single value
         Values to insert into `arr`.
     side : {'left', 'right'}, optional
         If 'left', the index of the first suitable location found is given.
@@ -1573,9 +1580,10 @@ def searchsorted(arr, value, side="left", sorter=None) -> np.ndarray:
             dtype = value_arr.dtype
 
         if is_scalar(value):
-            value = dtype.type(value)
+            # We know that value is int
+            value = cast(int, dtype.type(value))
         else:
-            value = pd_array(value, dtype=dtype)
+            value = pd_array(cast(ArrayLike, value), dtype=dtype)
     elif not (
         is_object_dtype(arr) or is_numeric_dtype(arr) or is_categorical_dtype(arr)
     ):

pandas/core/arrays/_mixins.py

@@ -2,6 +2,7 @@
 
 from functools import wraps
 from typing import (
+    TYPE_CHECKING,
     Any,
     Sequence,
     TypeVar,
@@ -13,7 +14,9 @@
 from pandas._libs import lib
 from pandas._libs.arrays import NDArrayBacked
 from pandas._typing import (
+    ArrayLike,
     F,
+    NumpySorter,
     PositionalIndexer2D,
     Shape,
     type_t,
@@ -45,6 +48,9 @@
     "NDArrayBackedExtensionArrayT", bound="NDArrayBackedExtensionArray"
 )
 
+if TYPE_CHECKING:
+    from typing import Literal
+
 
 def ravel_compat(meth: F) -> F:
     """
@@ -176,9 +182,14 @@ def _concat_same_type(
         return to_concat[0]._from_backing_data(new_values)  # type: ignore[arg-type]
 
     @doc(ExtensionArray.searchsorted)
-    def searchsorted(self, value, side="left", sorter=None):
-        value = self._validate_searchsorted_value(value)
-        return self._ndarray.searchsorted(value, side=side, sorter=sorter)
+    def searchsorted(
+        self,
+        value: ArrayLike | object,
+        side: Literal["left", "right"] = "left",
+        sorter: NumpySorter = None,
+    ) -> np.ndarray:
+        npvalue: np.ndarray = cast(np.ndarray, self._validate_searchsorted_value(value))
+        return self._ndarray.searchsorted(npvalue, side=side, sorter=sorter)
 
     def _validate_searchsorted_value(self, value):
         return value

pandas/core/arrays/base.py

@@ -26,6 +26,7 @@
     ArrayLike,
     Dtype,
     FillnaOptions,
+    NumpySorter,
     PositionalIndexer,
     Shape,
 )
@@ -812,7 +813,12 @@ def unique(self: ExtensionArrayT) -> ExtensionArrayT:
         uniques = unique(self.astype(object))
         return self._from_sequence(uniques, dtype=self.dtype)
 
-    def searchsorted(self, value, side="left", sorter=None):
+    def searchsorted(
+        self,
+        value: ArrayLike | object,
+        side: Literal["left", "right"] = "left",
+        sorter: NumpySorter = None,
+    ) -> np.ndarray:
         """
         Find indices where elements should be inserted to maintain order.
 
@@ -833,7 +839,7 @@ def searchsorted(self, value, side="left", sorter=None):
 
         Parameters
         ----------
-        value : array_like
+        value : array_like or a single value
             Values to insert into `self`.
         side : {'left', 'right'}, optional
             If 'left', the index of the first suitable location found is given.
@@ -1307,7 +1313,7 @@ def __hash__(self) -> int:
     # ------------------------------------------------------------------------
     # Non-Optimized Default Methods
 
-    def delete(self: ExtensionArrayT, loc) -> ExtensionArrayT:
+    def delete(self: ExtensionArrayT, loc: PositionalIndexer) -> ExtensionArrayT:
         indexer = np.delete(np.arange(len(self)), loc)
         return self.take(indexer)
 

pandas/core/arrays/period.py

@@ -7,6 +7,7 @@
     Any,
     Callable,
     Sequence,
+    cast,
 )
 
 import numpy as np
@@ -39,8 +40,10 @@
 )
 from pandas._typing import (
     AnyArrayLike,
+    ArrayLike,
     Dtype,
     NpDtype,
+    NumpySorter,
 )
 from pandas.util._decorators import (
     cache_readonly,
@@ -74,6 +77,8 @@
 import pandas.core.common as com
 
 if TYPE_CHECKING:
+    from typing import Literal
+
     from pandas.core.arrays import DatetimeArray
 
 _shared_doc_kwargs = {
@@ -642,12 +647,19 @@ def astype(self, dtype, copy: bool = True):
             return self.asfreq(dtype.freq)
         return super().astype(dtype, copy=copy)
 
-    def searchsorted(self, value, side="left", sorter=None) -> np.ndarray:
-        value = self._validate_searchsorted_value(value).view("M8[ns]")
+    def searchsorted(
+        self,
+        value: ArrayLike | object,
+        side: Literal["left", "right"] = "left",
+        sorter: NumpySorter = None,
+    ) -> np.ndarray:
+        npvalue = cast(
+            np.ndarray, self._validate_searchsorted_value(value).view("M8[ns]")
+        )
 
         # Cast to M8 to get datetime-like NaT placement
         m8arr = self._ndarray.view("M8[ns]")
-        return m8arr.searchsorted(value, side=side, sorter=sorter)
+        return m8arr.searchsorted(npvalue, side=side, sorter=sorter)
 
     def fillna(self, value=None, method=None, limit=None) -> PeriodArray:
         if method is not None:

pandas/core/arrays/sparse/array.py

@@ -7,6 +7,7 @@
 import numbers
 import operator
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Sequence,
@@ -25,8 +26,10 @@
 )
 from pandas._libs.tslibs import NaT
 from pandas._typing import (
+    ArrayLike,
     Dtype,
     NpDtype,
+    NumpySorter,
     Scalar,
 )
 from pandas.compat.numpy import function as nv
@@ -77,6 +80,9 @@
 
 import pandas.io.formats.printing as printing
 
+if TYPE_CHECKING:
+    from typing import Literal
+
 # ----------------------------------------------------------------------------
 # Array
 
@@ -996,7 +1002,13 @@ def _take_without_fill(self, indices) -> np.ndarray | SparseArray:
 
         return taken
 
-    def searchsorted(self, v, side="left", sorter=None):
+    def searchsorted(
+        self,
+        v: ArrayLike | object,
+        side: Literal["left", "right"] = "left",
+        sorter: NumpySorter = None,
+    ) -> np.ndarray:
+
         msg = "searchsorted requires high memory usage."
         warnings.warn(msg, PerformanceWarning, stacklevel=2)
         if not is_scalar(v):

pandas/core/base.py

@@ -23,6 +23,8 @@
     DtypeObj,
     FrameOrSeries,
     IndexLabel,
+    NumpySorter,
+    NumpyValueArrayLike,
     Shape,
     final,
 )
@@ -1226,7 +1228,12 @@ def factorize(self, sort: bool = False, na_sentinel: int | None = -1):
         """
 
     @doc(_shared_docs["searchsorted"], klass="Index")
-    def searchsorted(self, value, side="left", sorter=None) -> np.ndarray:
+    def searchsorted(
+        self,
+        value: NumpyValueArrayLike,
+        side: Literal["left", "right"] = "left",
+        sorter: NumpySorter = None,
+    ) -> np.ndarray:
         return algorithms.searchsorted(self._values, value, side=side, sorter=sorter)
 
     def drop_duplicates(self, keep="first"):

pandas/core/indexes/base.py

@@ -3564,7 +3564,7 @@ def _get_fill_indexer_searchsorted(
                 "if index and target are monotonic"
             )
 
-        side = "left" if method == "pad" else "right"
+        side: Literal["left", "right"] = "left" if method == "pad" else "right"
 
         # find exact matches first (this simplifies the algorithm)
         indexer = self.get_indexer(target)
@@ -5740,7 +5740,7 @@ def _maybe_cast_slice_bound(self, label, side: str_t, kind=no_default):
 
         return label
 
-    def _searchsorted_monotonic(self, label, side: str_t = "left"):
+    def _searchsorted_monotonic(self, label, side: Literal["left", "right"] = "left"):
         if self.is_monotonic_increasing:
             return self.searchsorted(label, side=side)
         elif self.is_monotonic_decreasing:
@@ -5754,7 +5754,7 @@ def _searchsorted_monotonic(self, label, side: str_t = "left"):
 
         raise ValueError("index must be monotonic increasing or decreasing")
 
-    def get_slice_bound(self, label, side: str_t, kind=None) -> int:
+    def get_slice_bound(self, label, side: Literal["left", "right"], kind=None) -> int:
         """
         Calculate slice bound that corresponds to given label.
 

pandas/core/indexes/extension.py

@@ -1,17 +1,21 @@
 """
 Shared methods for Index subclasses backed by ExtensionArray.
 """
+from __future__ import annotations
+
 from typing import (
+    TYPE_CHECKING,
     Hashable,
-    List,
-    Type,
     TypeVar,
-    Union,
 )
 
 import numpy as np
 
-from pandas._typing import ArrayLike
+from pandas._typing import (
+    ArrayLike,
+    NumpySorter,
+    Scalar,
+)
 from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import (
@@ -45,6 +49,9 @@
 from pandas.core.indexes.base import Index
 from pandas.core.ops import get_op_result_name
 
+if TYPE_CHECKING:
+    from typing import Literal
+
 _T = TypeVar("_T", bound="NDArrayBackedExtensionIndex")
 
 
@@ -120,7 +127,7 @@ def method(self, *args, **kwargs):
     return method
 
 
-def inherit_names(names: List[str], delegate, cache: bool = False, wrap: bool = False):
+def inherit_names(names: list[str], delegate, cache: bool = False, wrap: bool = False):
     """
     Class decorator to pin attributes from an ExtensionArray to a Index subclass.
 
@@ -230,20 +237,20 @@ class ExtensionIndex(Index):
     # The base class already passes through to _data:
     #  size, __len__, dtype
 
-    _data: Union[IntervalArray, NDArrayBackedExtensionArray]
+    _data: IntervalArray | NDArrayBackedExtensionArray
 
-    _data_cls: Union[
-        Type[Categorical],
-        Type[DatetimeArray],
-        Type[TimedeltaArray],
-        Type[PeriodArray],
-        Type[IntervalArray],
-    ]
+    _data_cls: (
+        type[Categorical]
+        | type[DatetimeArray]
+        | type[TimedeltaArray]
+        | type[PeriodArray]
+        | type[IntervalArray]
+    )
 
     @classmethod
     def _simple_new(
         cls,
-        array: Union[IntervalArray, NDArrayBackedExtensionArray],
+        array: IntervalArray | NDArrayBackedExtensionArray,
         name: Hashable = None,
     ):
         """
@@ -292,7 +299,12 @@ def __getitem__(self, key):
         deprecate_ndim_indexing(result)
         return result
 
-    def searchsorted(self, value, side="left", sorter=None) -> np.ndarray:
+    def searchsorted(
+        self,
+        value: ArrayLike | Scalar,
+        side: Literal["left", "right"] = "left",
+        sorter: NumpySorter = None,
+    ) -> np.ndarray:
         # overriding IndexOpsMixin improves performance GH#38083
         return self._data.searchsorted(value, side=side, sorter=sorter)