Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 34 additions & 14 deletions pandas-stubs/core/base.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ from typing import (
Any,
Generic,
Literal,
TypeAlias,
final,
overload,
)
Expand All @@ -22,18 +23,22 @@ from typing_extensions import Self

from pandas._typing import (
S1,
ArrayLike,
AxisIndex,
DropKeep,
DTypeLike,
GenericT,
GenericT_co,
NDFrameT,
Scalar,
SequenceNotStr,
SupportsDType,
np_1darray,
)
from pandas.util._decorators import cache_readonly

_ListLike: TypeAlias = ArrayLike | dict[str, np.ndarray] | SequenceNotStr[S1]
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not sure why this gets redefined in multiple files, but i've kept to the pattern

ok, or better to redefine in _typing?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you do want to use it here, then I think moving the definition to typing.pyi makes sense, but exercise care on the naming, because there is a ListLike in _typing that is used elsewhere and I'm not sure whether they should be the same or not. Maybe the name _ListLike is misleading.

Might be better to do the copying like you've done here, OR have series.pyi import _ListLike from core/base.pyi

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might be better to do the copying like you've done here

sure, keeping like this for now then, thanks


class NoNewAttributesMixin:
def __setattr__(self, key: str, value: Any) -> None: ...

Expand All @@ -51,7 +56,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
@property
def T(self) -> Self: ...
@property
def shape(self) -> tuple: ...
def shape(self) -> tuple[int, ...]: ...
@property
def ndim(self) -> int: ...
def item(self) -> S1: ...
Expand All @@ -67,41 +72,45 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
dtype: None = None,
copy: bool = False,
na_value: Scalar = ...,
**kwargs,
**kwargs: Any,
) -> np_1darray[GenericT_co]: ...
@overload
def to_numpy(
self,
dtype: np.dtype[GenericT] | SupportsDType[GenericT] | type[GenericT],
copy: bool = False,
na_value: Scalar = ...,
**kwargs,
**kwargs: Any,
) -> np_1darray[GenericT]: ...
@overload
def to_numpy(
self,
dtype: DTypeLike,
copy: bool = False,
na_value: Scalar = ...,
**kwargs,
**kwargs: Any,
) -> np_1darray: ...
@property
def empty(self) -> bool: ...
def max(self, axis=..., skipna: bool = ..., **kwargs): ...
def min(self, axis=..., skipna: bool = ..., **kwargs): ...
def max(
self, axis: AxisIndex | None = ..., skipna: bool = ..., **kwargs: Any
) -> S1: ...
def min(
self, axis: AxisIndex | None = ..., skipna: bool = ..., **kwargs: Any
) -> S1: ...
def argmax(
self,
axis: AxisIndex | None = ...,
skipna: bool = True,
*args,
**kwargs,
*args: Any,
**kwargs: Any,
) -> np.int64: ...
def argmin(
self,
axis: AxisIndex | None = ...,
skipna: bool = True,
*args,
**kwargs,
*args: Any,
**kwargs: Any,
) -> np.int64: ...
def tolist(self) -> list[S1]: ...
def to_list(self) -> list[S1]: ...
Expand All @@ -114,7 +123,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
normalize: Literal[False] = ...,
sort: bool = ...,
ascending: bool = ...,
bins=...,
bins: int | None = ...,
dropna: bool = ...,
) -> Series[int]: ...
@overload
Expand All @@ -123,7 +132,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
normalize: Literal[True],
sort: bool = ...,
ascending: bool = ...,
bins=...,
bins: int | None = ...,
dropna: bool = ...,
) -> Series[float]: ...
def nunique(self, dropna: bool = True) -> int: ...
Expand All @@ -136,7 +145,18 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
def factorize(
self, sort: bool = False, use_na_sentinel: bool = True
) -> tuple[np_1darray, np_1darray | Index | Categorical]: ...
@overload
def searchsorted(
self,
value: _ListLike,
side: Literal["left", "right"] = ...,
sorter: _ListLike | None = ...,
) -> np_1darray[np.intp]: ...
@overload
def searchsorted(
self, value, side: Literal["left", "right"] = ..., sorter=...
) -> int | list[int]: ...
self,
value: Scalar,
side: Literal["left", "right"] = ...,
sorter: _ListLike | None = ...,
) -> np.intp: ...
def drop_duplicates(self, *, keep: DropKeep = ...) -> Self: ...
4 changes: 2 additions & 2 deletions pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -851,14 +851,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
value: _ListLike,
side: Literal["left", "right"] = ...,
sorter: _ListLike | None = ...,
) -> list[int]: ...
) -> np_1darray[np.intp]: ...
@overload
def searchsorted(
self,
value: Scalar,
side: Literal["left", "right"] = ...,
sorter: _ListLike | None = ...,
) -> int: ...
) -> np.intp: ...
@overload
def compare(
self,
Expand Down
14 changes: 11 additions & 3 deletions tests/indexes/test_indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,6 @@
assert_type,
)

if TYPE_CHECKING:
from tests import Dtype # noqa: F401

from tests import (
PD_LTE_23,
TYPE_CHECKING_INVALID_USAGE,
Expand All @@ -30,6 +27,9 @@
pytest_warns_bounded,
)

if TYPE_CHECKING:
from tests import Dtype # noqa: F401


def test_index_unique() -> None:
df = pd.DataFrame({"x": [1, 2, 3, 4]}, index=pd.Index([1, 2, 3, 2]))
Expand Down Expand Up @@ -1489,6 +1489,14 @@ def test_index_naming() -> None:
check(assert_type(df.index.names, list[Hashable | None]), list)


def test_index_searchsorted() -> None:
idx = pd.Index([1, 2, 3])
check(assert_type(idx.searchsorted(1), np.intp), np.intp)
check(assert_type(idx.searchsorted([1]), "np_1darray[np.intp]"), np_1darray)
check(assert_type(idx.searchsorted(1, side="left"), np.intp), np.intp)
check(assert_type(idx.searchsorted(1, sorter=[1, 0, 2]), np.intp), np.intp)


def test_period_index_constructor() -> None:
check(
assert_type(pd.PeriodIndex(["2000"], dtype="period[D]"), pd.PeriodIndex),
Expand Down
Loading