TYP: ExtensionArray delete() and searchsorted()

pandas/_typing.py

@@ -91,6 +91,12 @@
 PandasScalar = Union["Period", "Timestamp", "Timedelta", "Interval"]
 Scalar = Union[PythonScalar, PandasScalar]
 
+# numpy compatible types
+NumpyArrayLike = Union[np.ndarray, Sequence[PythonScalar]]
+NumpyValueArrayLike = Union[PythonScalar, Sequence[PythonScalar]]
+NumpySorter = Union[int, bool, np.ndarray, Sequence[int], 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,
 )
@@ -813,7 +814,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.
 
@@ -834,7 +840,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.
@@ -1308,7 +1314,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,
 )
@@ -1225,7 +1227,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

@@ -3545,7 +3545,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)
@@ -5721,7 +5721,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:
@@ -5735,7 +5735,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: 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)