type IndexOpsMixin (#1357)

Author: MarcoGorelli

pandas-stubs/core/base.pyi

@@ -6,6 +6,7 @@ from typing import (
     Any,
     Generic,
     Literal,
+    TypeAlias,
     final,
     overload,
 )
@@ -22,18 +23,22 @@ from typing_extensions import Self
 
 from pandas._typing import (
     S1,
+    ArrayLike,
     AxisIndex,
     DropKeep,
     DTypeLike,
     GenericT,
     GenericT_co,
     NDFrameT,
     Scalar,
+    SequenceNotStr,
     SupportsDType,
     np_1darray,
 )
 from pandas.util._decorators import cache_readonly
 
+_ListLike: TypeAlias = ArrayLike | dict[str, np.ndarray] | SequenceNotStr[S1]
+
 class NoNewAttributesMixin:
     def __setattr__(self, key: str, value: Any) -> None: ...
 
@@ -51,7 +56,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
     @property
     def T(self) -> Self: ...
     @property
-    def shape(self) -> tuple: ...
+    def shape(self) -> tuple[int, ...]: ...
     @property
     def ndim(self) -> int: ...
     def item(self) -> S1: ...
@@ -67,41 +72,45 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
         dtype: None = None,
         copy: bool = False,
         na_value: Scalar = ...,
-        **kwargs,
+        **kwargs: Any,
     ) -> np_1darray[GenericT_co]: ...
     @overload
     def to_numpy(
         self,
         dtype: np.dtype[GenericT] | SupportsDType[GenericT] | type[GenericT],
         copy: bool = False,
         na_value: Scalar = ...,
-        **kwargs,
+        **kwargs: Any,
     ) -> np_1darray[GenericT]: ...
     @overload
     def to_numpy(
         self,
         dtype: DTypeLike,
         copy: bool = False,
         na_value: Scalar = ...,
-        **kwargs,
+        **kwargs: Any,
     ) -> np_1darray: ...
     @property
     def empty(self) -> bool: ...
-    def max(self, axis=..., skipna: bool = ..., **kwargs): ...
-    def min(self, axis=..., skipna: bool = ..., **kwargs): ...
+    def max(
+        self, axis: AxisIndex | None = ..., skipna: bool = ..., **kwargs: Any
+    ) -> S1: ...
+    def min(
+        self, axis: AxisIndex | None = ..., skipna: bool = ..., **kwargs: Any
+    ) -> S1: ...
     def argmax(
         self,
         axis: AxisIndex | None = ...,
         skipna: bool = True,
-        *args,
-        **kwargs,
+        *args: Any,
+        **kwargs: Any,
     ) -> np.int64: ...
     def argmin(
         self,
         axis: AxisIndex | None = ...,
         skipna: bool = True,
-        *args,
-        **kwargs,
+        *args: Any,
+        **kwargs: Any,
     ) -> np.int64: ...
     def tolist(self) -> list[S1]: ...
     def to_list(self) -> list[S1]: ...
@@ -114,7 +123,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
         normalize: Literal[False] = ...,
         sort: bool = ...,
         ascending: bool = ...,
-        bins=...,
+        bins: int | None = ...,
         dropna: bool = ...,
     ) -> Series[int]: ...
     @overload
@@ -123,7 +132,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
         normalize: Literal[True],
         sort: bool = ...,
         ascending: bool = ...,
-        bins=...,
+        bins: int | None = ...,
         dropna: bool = ...,
     ) -> Series[float]: ...
     def nunique(self, dropna: bool = True) -> int: ...
@@ -136,7 +145,18 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
     def factorize(
         self, sort: bool = False, use_na_sentinel: bool = True
     ) -> tuple[np_1darray, np_1darray | Index | Categorical]: ...
+    @overload
+    def searchsorted(
+        self,
+        value: _ListLike,
+        side: Literal["left", "right"] = ...,
+        sorter: _ListLike | None = ...,
+    ) -> np_1darray[np.intp]: ...
+    @overload
     def searchsorted(
-        self, value, side: Literal["left", "right"] = ..., sorter=...
-    ) -> int | list[int]: ...
+        self,
+        value: Scalar,
+        side: Literal["left", "right"] = ...,
+        sorter: _ListLike | None = ...,
+    ) -> np.intp: ...
     def drop_duplicates(self, *, keep: DropKeep = ...) -> Self: ...

pandas-stubs/core/series.pyi

@@ -851,14 +851,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
         value: _ListLike,
         side: Literal["left", "right"] = ...,
         sorter: _ListLike | None = ...,
-    ) -> list[int]: ...
+    ) -> np_1darray[np.intp]: ...
     @overload
     def searchsorted(
         self,
         value: Scalar,
         side: Literal["left", "right"] = ...,
         sorter: _ListLike | None = ...,
-    ) -> int: ...
+    ) -> np.intp: ...
     @overload
     def compare(
         self,

tests/indexes/test_indexes.py

@@ -19,9 +19,6 @@
     assert_type,
 )
 
-if TYPE_CHECKING:
-    from tests import Dtype  # noqa: F401
-
 from tests import (
     PD_LTE_23,
     TYPE_CHECKING_INVALID_USAGE,
@@ -30,6 +27,9 @@
     pytest_warns_bounded,
 )
 
+if TYPE_CHECKING:
+    from tests import Dtype  # noqa: F401
+
 
 def test_index_unique() -> None:
     df = pd.DataFrame({"x": [1, 2, 3, 4]}, index=pd.Index([1, 2, 3, 2]))
@@ -1489,6 +1489,14 @@ def test_index_naming() -> None:
     check(assert_type(df.index.names, list[Hashable | None]), list)
 
 
+def test_index_searchsorted() -> None:
+    idx = pd.Index([1, 2, 3])
+    check(assert_type(idx.searchsorted(1), np.intp), np.intp)
+    check(assert_type(idx.searchsorted([1]), "np_1darray[np.intp]"), np_1darray)
+    check(assert_type(idx.searchsorted(1, side="left"), np.intp), np.intp)
+    check(assert_type(idx.searchsorted(1, sorter=[1, 0, 2]), np.intp), np.intp)
+
+
 def test_period_index_constructor() -> None:
     check(
         assert_type(pd.PeriodIndex(["2000"], dtype="period[D]"), pd.PeriodIndex),