Merge branch 'clean-iterable' of github.com:MarcoGorelli/pandas-stubs into clean-iterable

Author: MarcoGorelli

pandas-stubs/core/indexes/base.pyi

@@ -106,6 +106,7 @@ from pandas._typing import (
     PandasFloatDtypeArg,
     PyArrowFloatDtypeArg,
     ReindexMethod,
+    Renamer,
     S2_contra,
     Scalar,
     SequenceNotStr,
@@ -532,28 +533,48 @@ 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: np_ndarray_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: Iterable[Any], level: Level | None = None
     ) -> np_1darray_bool: ...
@@ -562,11 +583,11 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
         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[int | np.intp, int | np.intp]: ...
     def delete(
         self, loc: np.integer | int | AnyArrayLikeInt | Sequence[int]
     ) -> Self: ...

pandas-stubs/core/indexes/interval.pyi

@@ -214,9 +214,6 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin):
     @property
     def is_overlapping(self) -> bool: ...
     def get_loc(self, key: Label) -> int | slice | np_1darray_bool: ...
-    def get_indexer_non_unique(
-        self, target: Index
-    ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ...
     @property
     def left(self) -> Index: ...
     @property

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,
@@ -1590,3 +1592,77 @@ def test_index_setitem() -> None:
     idx = pd.Index([1, 2])
     if TYPE_CHECKING_INVALID_USAGE:
         idx[0] = 999  # type: ignore[index]  # pyright: ignore[reportIndexIssue]
+
+
+def test_index_putmask() -> None:
+    idx = pd.Index([1, 2])
+    check(assert_type(idx.putmask([True, False], 11.4), "pd.Index"), pd.Index)
+    check(assert_type(idx.putmask(np.array([True, False]), 11.4), "pd.Index"), pd.Index)
+    check(
+        assert_type(idx.putmask(pd.Series([True, False]), 11.4), "pd.Index"), pd.Index
+    )
+    check(assert_type(idx.putmask(pd.Index([True, False]), 11.4), "pd.Index"), pd.Index)
+    check(assert_type(idx.putmask(pd.array([True, False]), 11.5), "pd.Index"), pd.Index)
+
+
+def test_index_asof() -> None:
+    check(assert_type(pd.Index([1, 2]).asof(1), "Scalar"), np.integer)
+    check(assert_type(pd.Index(["a", "b", "c"]).asof("c"), "Scalar"), str)
+
+
+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_index, indexer = idx.sort_values(return_indexer=True)
+    check(assert_type(sorted_index, pd.DatetimeIndex), pd.DatetimeIndex)
+    check(assert_type(indexer, np_1darray_intp), np_1darray_intp)
+
+
+def test_index_get_indexer_non_unique() -> None:
+    idx = pd.Index([1, 3])
+    indexer, missing = idx.get_indexer_non_unique(pd.Index([3]))
+    check(assert_type(indexer, np_1darray_intp), np_1darray_intp)
+    check(assert_type(missing, 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])
+    start, end = idx.slice_locs(0, 1)
+    check(assert_type(start, np.intp | int), np.integer)
+    check(assert_type(end, np.intp | int), int)