diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 584548d97..6223a242a 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -99,6 +99,7 @@ from pandas._typing import ( MaskType, NaPosition, ReindexMethod, + Renamer, S2_contra, Scalar, SequenceNotStr, @@ -513,39 +514,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: 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, 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[int | np.intp, int | np.intp]: ... def delete( self, loc: np.integer | int | AnyArrayLikeInt | Sequence[int] ) -> Self: ... diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 0b4d18517..186a7a0a1 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/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 diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index f155efaff..f876c19aa 100644 --- a/tests/indexes/test_indexes.py +++ b/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, @@ -1582,3 +1584,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)