From ba7d1452df6bd93dc3658dc7e6b868b62b13d454 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 25 Nov 2025 16:46:11 +0000 Subject: [PATCH 1/9] type Index methods: putmask, asof, asof_locs, sort_values, get_indexer_non_unique, get_indexer_for, map, get_slice_bound --- pandas-stubs/core/indexes/base.pyi | 48 ++++++++++++++------ tests/indexes/test_indexes.py | 70 ++++++++++++++++++++++++++++++ 2 files changed, 105 insertions(+), 13 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 584548d97..021912263 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/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: ... diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index f155efaff..b101fa34c 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,71 @@ 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) + + +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) From bde64a18e5923163d36e78f55f63ed3b511016d4 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 25 Nov 2025 17:10:07 +0000 Subject: [PATCH 2/9] remove redundant `IntervalIndexer.get_indexer_non_unique` --- pandas-stubs/core/indexes/interval.pyi | 3 --- 1 file changed, 3 deletions(-) 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 From 209f11801f4ce49a75ab60bfcd7a8b5993d0d23e Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 25 Nov 2025 17:20:42 +0000 Subject: [PATCH 3/9] slice_locs can return either `np.integer` or `int` --- pandas-stubs/core/indexes/base.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 021912263..0dfbca5ef 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -567,7 +567,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): 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]: ... + ) -> tuple[int | np.integer, int | np.integer]: ... def delete( self, loc: np.integer | int | AnyArrayLikeInt | Sequence[int] ) -> Self: ... From a90ea69b559158b6e2477104f4d5933d48570e2e Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Tue, 25 Nov 2025 18:34:17 +0000 Subject: [PATCH 4/9] fix slice_locs test --- tests/indexes/test_indexes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index b101fa34c..a16c590bf 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1650,5 +1650,5 @@ def test_index_get_slice_bound() -> None: 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) + check(assert_type(res[0], np.integer | int), np.integer) + check(assert_type(res[1], np.integer | int), int) From 6439619a2fcac831924fda5d750890218bdc5c8a Mon Sep 17 00:00:00 2001 From: Marco Edward Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 26 Nov 2025 10:36:29 +0000 Subject: [PATCH 5/9] Apply suggestions from code review Co-authored-by: Yi-Fan Wang --- pandas-stubs/core/indexes/base.pyi | 2 +- tests/indexes/test_indexes.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 0dfbca5ef..96d8aaa49 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -526,7 +526,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): @final def asof(self, label: Scalar) -> Scalar: ... def asof_locs( - self, where: DatetimeIndex, mask: npt.NDArray[np.bool_] + self, where: DatetimeIndex, mask: np_ndarray_bool ) -> np_1darray_intp: ... @overload def sort_values( diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index a16c590bf..d795e92b1 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1612,9 +1612,9 @@ def test_index_asof_locs() -> None: 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) + 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: From 990ce3b696fe3a7bd85c6c96316d7d17b8692982 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 26 Nov 2025 10:39:26 +0000 Subject: [PATCH 6/9] np.integer -> np.intp --- pandas-stubs/core/indexes/base.pyi | 3 +-- tests/indexes/test_indexes.py | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 96d8aaa49..6223a242a 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -28,7 +28,6 @@ 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, @@ -567,7 +566,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): 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.integer, int | np.integer]: ... + ) -> tuple[int | np.intp, int | np.intp]: ... def delete( self, loc: np.integer | int | AnyArrayLikeInt | Sequence[int] ) -> Self: ... diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index d795e92b1..a44e0dca9 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1650,5 +1650,5 @@ def test_index_get_slice_bound() -> None: def test_index_slice_locs() -> None: idx = pd.Index([1, 3]) res = idx.slice_locs(0, 1) - check(assert_type(res[0], np.integer | int), np.integer) - check(assert_type(res[1], np.integer | int), int) + check(assert_type(res[0], np.intp | int), np.integer) + check(assert_type(res[1], np.intp | int), int) From 7424d2e662abd816c7df538b8177b98b8a79c367 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 26 Nov 2025 14:12:23 +0000 Subject: [PATCH 7/9] test more mask types for `putmask` --- tests/indexes/test_indexes.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index a44e0dca9..bcdd68bbc 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1589,6 +1589,12 @@ def test_index_setitem() -> None: 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: From d4a033fe1dcddfc0c8a8c6b2b24f43a323e24f36 Mon Sep 17 00:00:00 2001 From: Marco Edward Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 26 Nov 2025 15:08:35 +0000 Subject: [PATCH 8/9] Apply suggestions from code review Co-authored-by: Yi-Fan Wang --- tests/indexes/test_indexes.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index bcdd68bbc..145cd156e 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1625,9 +1625,9 @@ def test_index_sort_values() -> None: 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) + 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: @@ -1655,6 +1655,6 @@ def test_index_get_slice_bound() -> None: def test_index_slice_locs() -> None: idx = pd.Index([1, 3]) - res = idx.slice_locs(0, 1) - check(assert_type(res[0], np.intp | int), np.integer) - check(assert_type(res[1], np.intp | int), int) + start, end = idx.slice_locs(0, 1) + check(assert_type(start, np.intp | int), np.integer) + check(assert_type(end, np.intp | int), int) From bd2654b13354e2236a2822cbe83db3cfb69198eb Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 26 Nov 2025 16:23:16 +0000 Subject: [PATCH 9/9] test asof with str --- tests/indexes/test_indexes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index bcdd68bbc..13b5dfefc 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1598,8 +1598,8 @@ def test_index_putmask() -> None: def test_index_asof() -> None: - idx = pd.Index([1, 2]) - check(assert_type(idx.asof(1), "Scalar"), np.integer) + 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: