Skip to content
Merged
47 changes: 34 additions & 13 deletions pandas-stubs/core/indexes/base.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ from pandas._typing import (
MaskType,
NaPosition,
ReindexMethod,
Renamer,
S2_contra,
Scalar,
SequenceNotStr,
Expand Down Expand Up @@ -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: ...
Expand Down
3 changes: 0 additions & 3 deletions pandas-stubs/core/indexes/interval.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
76 changes: 76 additions & 0 deletions tests/indexes/test_indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
assert_type,
)

from pandas._typing import Scalar # noqa: F401

from tests import (
PD_LTE_23,
TYPE_CHECKING_INVALID_USAGE,
Expand Down Expand Up @@ -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)