Skip to content

Commit b82986c

Browse files
MarcoGorellicmp0xff
MarcoGorelli
and
cmp0xff
authored
type Index methods: putmask, asof, asof_locs, sort_values, get_indexer_non_unique, get_indexer_for, map, get_slice_bound (#1505)
* type Index methods: putmask, asof, asof_locs, sort_values, get_indexer_non_unique, get_indexer_for, map, get_slice_bound * remove redundant `IntervalIndexer.get_indexer_non_unique` * slice_locs can return either `np.integer` or `int` * fix slice_locs test * Apply suggestions from code review Co-authored-by: Yi-Fan Wang <cmp0xff@users.noreply.github.com> * np.integer -> np.intp * test more mask types for `putmask` * Apply suggestions from code review Co-authored-by: Yi-Fan Wang <cmp0xff@users.noreply.github.com> * test asof with str --------- Co-authored-by: Yi-Fan Wang <cmp0xff@users.noreply.github.com>
1 parent 4b1436f commit b82986c

File tree

3 files changed

+110
-16
lines changed

3 files changed

+110
-16
lines changed

‎pandas-stubs/core/indexes/base.pyi‎

Lines changed: 34 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ from pandas._typing import (
106106
PandasFloatDtypeArg,
107107
PyArrowFloatDtypeArg,
108108
ReindexMethod,
109+
Renamer,
109110
S2_contra,
110111
Scalar,
111112
SequenceNotStr,
@@ -532,39 +533,59 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
532533
) -> Index[C2]: ...
533534
@overload
534535
def append(self, other: Index | Sequence[Index]) -> Index: ...
535-
def putmask(self, mask, value): ...
536+
def putmask(
537+
self,
538+
mask: Sequence[bool] | np_ndarray_bool | BooleanArray | IndexOpsMixin[bool],
539+
value: Scalar,
540+
) -> Index: ...
536541
def equals(self, other: Any) -> bool: ...
537542
@final
538543
def identical(self, other: Any) -> bool: ...
539544
@final
540-
def asof(self, label): ...
541-
def asof_locs(self, where, mask): ...
545+
def asof(self, label: Scalar) -> Scalar: ...
546+
def asof_locs(
547+
self, where: DatetimeIndex, mask: np_ndarray_bool
548+
) -> np_1darray_intp: ...
549+
@overload
542550
def sort_values(
543551
self,
544552
*,
545-
return_indexer: bool = ...,
546-
ascending: bool = ...,
547-
na_position: NaPosition = ...,
553+
return_indexer: Literal[False] = False,
554+
ascending: bool = True,
555+
na_position: NaPosition = "last",
548556
key: Callable[[Index], Index] | None = None,
549-
): ...
557+
) -> Self: ...
558+
@overload
559+
def sort_values(
560+
self,
561+
*,
562+
return_indexer: Literal[True],
563+
ascending: bool = True,
564+
na_position: NaPosition = "last",
565+
key: Callable[[Index], Index] | None = None,
566+
) -> tuple[Self, np_1darray_intp]: ...
550567
@final
551568
def sort(self, *args: Any, **kwargs: Any) -> None: ...
552569
def argsort(self, *args: Any, **kwargs: Any) -> np_1darray_intp: ...
553-
def get_indexer_non_unique(self, target): ...
570+
def get_indexer_non_unique(
571+
self, target: Index
572+
) -> tuple[np_1darray_intp, np_1darray_intp]: ...
554573
@final
555-
def get_indexer_for(self, target, **kwargs: Any): ...
556-
def map(self, mapper, na_action=...) -> Index: ...
574+
def get_indexer_for(self, target: Index) -> np_1darray_intp: ...
575+
def map(
576+
self, mapper: Renamer, na_action: Literal["ignore"] | None = None
577+
) -> Index: ...
557578
def isin(self, values, level=...) -> np_1darray_bool: ...
558579
def slice_indexer(
559580
self,
560581
start: Label | None = None,
561582
end: Label | None = None,
562583
step: int | None = None,
563-
): ...
564-
def get_slice_bound(self, label, side): ...
584+
) -> slice: ...
585+
def get_slice_bound(self, label: Scalar, side: Literal["left", "right"]) -> int: ...
565586
def slice_locs(
566587
self, start: SliceType = None, end: SliceType = None, step: int | None = None
567-
): ...
588+
) -> tuple[int | np.intp, int | np.intp]: ...
568589
def delete(
569590
self, loc: np.integer | int | AnyArrayLikeInt | Sequence[int]
570591
) -> Self: ...

‎pandas-stubs/core/indexes/interval.pyi‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -214,9 +214,6 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin):
214214
@property
215215
def is_overlapping(self) -> bool: ...
216216
def get_loc(self, key: Label) -> int | slice | np_1darray_bool: ...
217-
def get_indexer_non_unique(
218-
self, target: Index
219-
) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ...
220217
@property
221218
def left(self) -> Index: ...
222219
@property

‎tests/indexes/test_indexes.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
assert_type,
2424
)
2525

26+
from pandas._typing import Scalar # noqa: F401
27+
2628
from tests import (
2729
PD_LTE_23,
2830
TYPE_CHECKING_INVALID_USAGE,
@@ -1582,3 +1584,77 @@ def test_index_setitem() -> None:
15821584
idx = pd.Index([1, 2])
15831585
if TYPE_CHECKING_INVALID_USAGE:
15841586
idx[0] = 999 # type: ignore[index] # pyright: ignore[reportIndexIssue]
1587+
1588+
1589+
def test_index_putmask() -> None:
1590+
idx = pd.Index([1, 2])
1591+
check(assert_type(idx.putmask([True, False], 11.4), "pd.Index"), pd.Index)
1592+
check(assert_type(idx.putmask(np.array([True, False]), 11.4), "pd.Index"), pd.Index)
1593+
check(
1594+
assert_type(idx.putmask(pd.Series([True, False]), 11.4), "pd.Index"), pd.Index
1595+
)
1596+
check(assert_type(idx.putmask(pd.Index([True, False]), 11.4), "pd.Index"), pd.Index)
1597+
check(assert_type(idx.putmask(pd.array([True, False]), 11.5), "pd.Index"), pd.Index)
1598+
1599+
1600+
def test_index_asof() -> None:
1601+
check(assert_type(pd.Index([1, 2]).asof(1), "Scalar"), np.integer)
1602+
check(assert_type(pd.Index(["a", "b", "c"]).asof("c"), "Scalar"), str)
1603+
1604+
1605+
def test_index_asof_locs() -> None:
1606+
idx = pd.DatetimeIndex(["2020-01-01", "2020-01-02", "2020-01-03"])
1607+
check(
1608+
assert_type(
1609+
idx.asof_locs(
1610+
pd.DatetimeIndex(["2020-01-01 11:00"]), np.array([True, True, True])
1611+
),
1612+
np_1darray_intp,
1613+
),
1614+
np_1darray_intp,
1615+
)
1616+
1617+
1618+
def test_index_sort_values() -> None:
1619+
idx = pd.DatetimeIndex(["2020-01-01", "2020-01-02", "2020-01-03"])
1620+
check(assert_type(idx.sort_values(), pd.DatetimeIndex), pd.DatetimeIndex)
1621+
sorted_index, indexer = idx.sort_values(return_indexer=True)
1622+
check(assert_type(sorted_index, pd.DatetimeIndex), pd.DatetimeIndex)
1623+
check(assert_type(indexer, np_1darray_intp), np_1darray_intp)
1624+
1625+
1626+
def test_index_get_indexer_non_unique() -> None:
1627+
idx = pd.Index([1, 3])
1628+
indexer, missing = idx.get_indexer_non_unique(pd.Index([3]))
1629+
check(assert_type(indexer, np_1darray_intp), np_1darray_intp)
1630+
check(assert_type(missing, np_1darray_intp), np_1darray_intp)
1631+
1632+
1633+
def test_index_get_indexer_for() -> None:
1634+
idx = pd.Index([1, 3])
1635+
check(
1636+
assert_type(idx.get_indexer_for(pd.Index([3])), np_1darray_intp),
1637+
np_1darray_intp,
1638+
)
1639+
1640+
1641+
def test_index_map() -> None:
1642+
idx = pd.Index([1, 3])
1643+
check(assert_type(idx.map(lambda x: str(x)), pd.Index), pd.Index)
1644+
1645+
1646+
def test_index_slice_indexer() -> None:
1647+
idx = pd.Index([1, 3])
1648+
check(assert_type(idx.slice_indexer(0, 1), slice), slice)
1649+
1650+
1651+
def test_index_get_slice_bound() -> None:
1652+
idx = pd.Index([1, 3])
1653+
check(assert_type(idx.get_slice_bound(1, side="left"), int), int)
1654+
1655+
1656+
def test_index_slice_locs() -> None:
1657+
idx = pd.Index([1, 3])
1658+
start, end = idx.slice_locs(0, 1)
1659+
check(assert_type(start, np.intp | int), np.integer)
1660+
check(assert_type(end, np.intp | int), int)

0 commit comments

Comments
 (0)