type Index methods: putmask, asof, asof_locs, sort_values, get_indexer_non_unique, get_indexer_for, map, get_slice_bound

Author: MarcoGorelli

pandas-stubs/core/indexes/base.pyi

@@ -28,6 +28,7 @@ from _typeshed import (
     _T_contra,
 )
 import numpy as np
+import numpy.typing as npt
 from pandas.core.arrays.boolean import BooleanArray
 from pandas.core.base import (
     ArrayIndexTimedeltaNoSeq,
@@ -99,6 +100,7 @@ from pandas._typing import (
     MaskType,
     NaPosition,
     ReindexMethod,
+    Renamer,
     S2_contra,
     Scalar,
     SequenceNotStr,
@@ -513,39 +515,59 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     ) -> Index[C2]: ...
     @overload
     def append(self, other: Index | Sequence[Index]) -> Index: ...
-    def putmask(self, mask, value): ...
+    def putmask(
+        self,
+        mask: Sequence[bool] | np_ndarray_bool | BooleanArray | IndexOpsMixin[bool],
+        value: Scalar,
+    ) -> Index: ...
     def equals(self, other: Any) -> bool: ...
     @final
     def identical(self, other: Any) -> bool: ...
     @final
-    def asof(self, label): ...
-    def asof_locs(self, where, mask): ...
+    def asof(self, label: Scalar) -> Scalar: ...
+    def asof_locs(
+        self, where: DatetimeIndex, mask: npt.NDArray[np.bool_]
+    ) -> np_1darray_intp: ...
+    @overload
     def sort_values(
         self,
         *,
-        return_indexer: bool = ...,
-        ascending: bool = ...,
-        na_position: NaPosition = ...,
+        return_indexer: Literal[False] = False,
+        ascending: bool = True,
+        na_position: NaPosition = "last",
         key: Callable[[Index], Index] | None = None,
-    ): ...
+    ) -> Self: ...
+    @overload
+    def sort_values(
+        self,
+        *,
+        return_indexer: Literal[True],
+        ascending: bool = True,
+        na_position: NaPosition = "last",
+        key: Callable[[Index], Index] | None = None,
+    ) -> tuple[Self, np_1darray_intp]: ...
     @final
     def sort(self, *args: Any, **kwargs: Any) -> None: ...
     def argsort(self, *args: Any, **kwargs: Any) -> np_1darray_intp: ...
-    def get_indexer_non_unique(self, target): ...
+    def get_indexer_non_unique(
+        self, target: Index
+    ) -> tuple[np_1darray_intp, np_1darray_intp]: ...
     @final
-    def get_indexer_for(self, target, **kwargs: Any): ...
-    def map(self, mapper, na_action=...) -> Index: ...
+    def get_indexer_for(self, target: Index) -> np_1darray_intp: ...
+    def map(
+        self, mapper: Renamer, na_action: Literal["ignore"] | None = None
+    ) -> Index: ...
     def isin(self, values, level=...) -> np_1darray_bool: ...
     def slice_indexer(
         self,
         start: Label | None = None,
         end: Label | None = None,
         step: int | None = None,
-    ): ...
-    def get_slice_bound(self, label, side): ...
+    ) -> slice: ...
+    def get_slice_bound(self, label: Scalar, side: Literal["left", "right"]) -> int: ...
     def slice_locs(
         self, start: SliceType = None, end: SliceType = None, step: int | None = None
-    ): ...
+    ) -> tuple[np.integer, int]: ...
     def delete(
         self, loc: np.integer | int | AnyArrayLikeInt | Sequence[int]
     ) -> Self: ...

tests/indexes/test_indexes.py

@@ -23,6 +23,8 @@
     assert_type,
 )
 
+from pandas._typing import Scalar  # noqa: F401
+
 from tests import (
     PD_LTE_23,
     TYPE_CHECKING_INVALID_USAGE,
@@ -1578,7 +1580,69 @@ def test_index_droplevel() -> None:
     check(assert_type(mi.droplevel(0), pd.MultiIndex | pd.Index), pd.Index)
 
 
-def test_index_setitem() -> None:
+def test_index_putmask() -> None:
     idx = pd.Index([1, 2])
-    if TYPE_CHECKING_INVALID_USAGE:
-        idx[0] = 999  # type: ignore[index]  # pyright: ignore[reportIndexIssue]
+    check(assert_type(idx.putmask([True, False], 11.4), "pd.Index"), pd.Index)
+
+
+def test_index_asof() -> None:
+    idx = pd.Index([1, 2])
+    check(assert_type(idx.asof(1), "Scalar"), np.integer)
+
+
+def test_index_asof_locs() -> None:
+    idx = pd.DatetimeIndex(["2020-01-01", "2020-01-02", "2020-01-03"])
+    check(
+        assert_type(
+            idx.asof_locs(
+                pd.DatetimeIndex(["2020-01-01 11:00"]), np.array([True, True, True])
+            ),
+            np_1darray_intp,
+        ),
+        np_1darray_intp,
+    )
+
+
+def test_index_sort_values() -> None:
+    idx = pd.DatetimeIndex(["2020-01-01", "2020-01-02", "2020-01-03"])
+    check(assert_type(idx.sort_values(), pd.DatetimeIndex), pd.DatetimeIndex)
+    sorted_ = idx.sort_values(return_indexer=True)
+    check(assert_type(sorted_[0], pd.DatetimeIndex), pd.DatetimeIndex)
+    check(assert_type(sorted_[1], np_1darray_intp), np_1darray_intp)
+
+
+def test_index_get_indexer_non_unique() -> None:
+    idx = pd.Index([1, 3])
+    res = idx.get_indexer_non_unique(pd.Index([3]))
+    check(assert_type(res[0], np_1darray_intp), np_1darray_intp)
+    check(assert_type(res[1], np_1darray_intp), np_1darray_intp)
+
+
+def test_index_get_indexer_for() -> None:
+    idx = pd.Index([1, 3])
+    check(
+        assert_type(idx.get_indexer_for(pd.Index([3])), np_1darray_intp),
+        np_1darray_intp,
+    )
+
+
+def test_index_map() -> None:
+    idx = pd.Index([1, 3])
+    check(assert_type(idx.map(lambda x: str(x)), pd.Index), pd.Index)
+
+
+def test_index_slice_indexer() -> None:
+    idx = pd.Index([1, 3])
+    check(assert_type(idx.slice_indexer(0, 1), slice), slice)
+
+
+def test_index_get_slice_bound() -> None:
+    idx = pd.Index([1, 3])
+    check(assert_type(idx.get_slice_bound(1, side="left"), int), int)
+
+
+def test_index_slice_locs() -> None:
+    idx = pd.Index([1, 3])
+    res = idx.slice_locs(0, 1)
+    check(assert_type(res[0], np.integer), np.integer)
+    check(assert_type(res[1], int), int)