Merge branch 'main' of github.com:loicdiridollou/pandas-stubs into gh1484_ty_29a

Author: loicdiridollou

pandas-stubs/core/base.pyi

@@ -10,6 +10,7 @@ from typing import (
     Literal,
     Protocol,
     TypeAlias,
+    TypeVar,
     final,
     overload,
     type_check_only,
@@ -53,6 +54,8 @@ from pandas._typing import (
 )
 from pandas.util._decorators import cache_readonly
 
+T_INTERVAL_NP = TypeVar("T_INTERVAL_NP", bound=np.bytes_ | np.str_)
+
 class NoNewAttributesMixin:
     def __setattr__(self, key: str, value: Any) -> None: ...
 

pandas-stubs/core/indexes/base.pyi

@@ -32,6 +32,7 @@ import numpy as np
 from pandas.core.arrays.boolean import BooleanArray
 from pandas.core.arrays.floating import FloatingArray
 from pandas.core.base import (
+    T_INTERVAL_NP,
     ArrayIndexTimedeltaNoSeq,
     ElementOpsMixin,
     IndexComplex,
@@ -106,6 +107,7 @@ from pandas._typing import (
     PandasFloatDtypeArg,
     PyArrowFloatDtypeArg,
     ReindexMethod,
+    Renamer,
     S2_contra,
     Scalar,
     SequenceNotStr,
@@ -532,39 +534,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: Literal[False] = False,
+        ascending: bool = True,
+        na_position: NaPosition = "last",
+        key: Callable[[Index], Index] | None = None,
+    ) -> Self: ...
+    @overload
     def sort_values(
         self,
         *,
-        return_indexer: bool = ...,
-        ascending: bool = ...,
-        na_position: NaPosition = ...,
+        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: ...
@@ -1168,6 +1190,14 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
 
 @type_check_only
 class _IndexSubclassBase(Index[S1], Generic[S1, GenericT_co]):
+    @overload
+    def to_numpy(
+        self: _IndexSubclassBase[Interval],
+        dtype: type[T_INTERVAL_NP],
+        copy: bool = False,
+        na_value: Scalar = ...,
+        **kwargs: Any,
+    ) -> np_1darray: ...
     @overload
     def to_numpy(
         self,
@@ -1185,7 +1215,7 @@ class _IndexSubclassBase(Index[S1], Generic[S1, GenericT_co]):
         **kwargs: Any,
     ) -> np_1darray[GenericT]: ...
     @overload
-    def to_numpy(
+    def to_numpy(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         dtype: DTypeLike,
         copy: bool = False,

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

pandas-stubs/core/series.pyi

@@ -29,7 +29,6 @@ from typing import (
     NoReturn,
     Protocol,
     TypeAlias,
-    TypeVar,
     final,
     overload,
     type_check_only,
@@ -71,6 +70,7 @@ from pandas.core.arrays.datetimes import DatetimeArray
 from pandas.core.arrays.floating import FloatingArray
 from pandas.core.arrays.timedeltas import TimedeltaArray
 from pandas.core.base import (
+    T_INTERVAL_NP,
     ArrayIndexSeriesTimedeltaNoSeq,
     ArrayIndexTimedeltaNoSeq,
     ElementOpsMixin,
@@ -189,6 +189,7 @@ from pandas._typing import (
     MaskType,
     NaPosition,
     NsmallestNlargestKeep,
+    NumpyStrDtypeArg,
     ObjectDtypeArg,
     PandasAstypeComplexDtypeArg,
     PandasAstypeFloatDtypeArg,
@@ -251,8 +252,6 @@ from pandas.core.dtypes.dtypes import CategoricalDtype
 
 from pandas.plotting import PlotAccessor
 
-_T_INTERVAL_NP = TypeVar("_T_INTERVAL_NP", bound=np.bytes_ | np.str_)
-
 @type_check_only
 class _SupportsAdd(Protocol[_T_co]):
     def __add__(self, value: Self, /) -> _T_co: ...
@@ -4507,7 +4506,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
         copy: bool = False,
         na_value: Scalar = ...,
         **kwargs: Any,
-    ) -> np_1darray_bytes: ...
+    ) -> np_1darray: ...
     @overload
     def to_numpy(
         self: Series[Interval],
@@ -4519,11 +4518,11 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
     @overload
     def to_numpy(
         self: Series[Interval],
-        dtype: type[_T_INTERVAL_NP],
+        dtype: type[T_INTERVAL_NP],
         copy: bool = False,
         na_value: Scalar = ...,
         **kwargs: Any,
-    ) -> np_1darray[_T_INTERVAL_NP]: ...
+    ) -> np_1darray: ...
     @overload
     def to_numpy(
         self: Series[int],
@@ -4559,12 +4558,28 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
     @overload
     def to_numpy(
         self: Series[_str],
-        dtype: DTypeLike | None = None,
+        dtype: NumpyStrDtypeArg,
         copy: bool = False,
         na_value: Scalar = ...,
         **kwargs: Any,
     ) -> np_1darray_str: ...
     @overload
+    def to_numpy(
+        self: Series[_str],
+        dtype: DTypeLike,
+        copy: bool = False,
+        na_value: Scalar = ...,
+        **kwargs: Any,
+    ) -> np_1darray: ...
+    @overload
+    def to_numpy(
+        self: Series[_str],
+        dtype: None = None,
+        copy: bool = False,
+        na_value: Scalar = ...,
+        **kwargs: Any,
+    ) -> np_1darray_object: ...
+    @overload
     def to_numpy(
         self: Series[bytes],
         dtype: DTypeLike | None = None,

tests/__init__.py

@@ -439,11 +439,11 @@
     _S = TypeVar("_S", bound=tuple[int, ...])
     # Separately define here so pytest works
     np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[_G]]
-    np_1darray_bool: TypeAlias = np.ndarray[tuple[int], np.bool_]
-    np_1darray_str: TypeAlias = np.ndarray[tuple[int], np.str_]
-    np_1darray_bytes: TypeAlias = np.ndarray[tuple[int], np.bytes_]
-    np_1darray_complex: TypeAlias = np.ndarray[tuple[int], np.complexfloating]
-    np_1darray_object: TypeAlias = np.ndarray[tuple[int], np.object_]
+    np_1darray_bool: TypeAlias = np_1darray[np.bool_]
+    np_1darray_str: TypeAlias = np_1darray[np.str_]
+    np_1darray_bytes: TypeAlias = np_1darray[np.bytes_]
+    np_1darray_complex: TypeAlias = np_1darray[np.complexfloating]
+    np_1darray_object: TypeAlias = np_1darray[np.object_]
     np_1darray_intp: TypeAlias = np_1darray[np.intp]
     np_1darray_int64: TypeAlias = np_1darray[np.int64]
     np_1darray_anyint: TypeAlias = np_1darray[np.integer]

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)

tests/series/test_series.py

@@ -73,7 +73,6 @@
     np_1darray_dt,
     np_1darray_float,
     np_1darray_object,
-    np_1darray_str,
     np_1darray_td,
     np_ndarray_num,
     pytest_warns_bounded,
@@ -2004,18 +2003,22 @@ def test_dtype_type() -> None:
 
 def test_types_to_numpy() -> None:
     s = pd.Series(["a", "b", "c"], dtype=str)
-    check(assert_type(s.to_numpy(), np_1darray_str), np_1darray_str)
+    check(assert_type(s.to_numpy(), np_1darray_object), np_1darray_object)
+    check(  # <U1, not str_
+        assert_type(s.to_numpy(dtype="str", copy=True), np_1darray), np_1darray
+    )
+    check(assert_type(s.to_numpy(na_value=0), np_1darray_object), np_1darray_object)
     check(
-        assert_type(s.to_numpy(dtype="str", copy=True), np_1darray_str), np_1darray_str
+        assert_type(s.to_numpy(na_value=np.int32(4)), np_1darray_object),
+        np_1darray_object,
     )
-    check(assert_type(s.to_numpy(na_value=0), np_1darray_str), np_1darray_str)
-    check(assert_type(s.to_numpy(na_value=np.int32(4)), np_1darray_str), np_1darray_str)
     check(
-        assert_type(s.to_numpy(na_value=np.float16(4)), np_1darray_str), np_1darray_str
+        assert_type(s.to_numpy(na_value=np.float16(4)), np_1darray_object),
+        np_1darray_object,
     )
     check(
-        assert_type(s.to_numpy(na_value=np.complex128(4, 7)), np_1darray_str),
-        np_1darray_str,
+        assert_type(s.to_numpy(na_value=np.complex128(4, 7)), np_1darray_object),
+        np_1darray_object,
     )
 
     check(assert_type(pd.Series().to_numpy(), np_1darray), np_1darray)
@@ -2024,7 +2027,7 @@ def test_types_to_numpy() -> None:
 def test_to_numpy() -> None:
     """Test Series.to_numpy for different types."""
     s_str = pd.Series(["a", "b", "c"], dtype=str)
-    check(assert_type(s_str.to_numpy(), np_1darray_str), np_1darray_str)
+    check(assert_type(s_str.to_numpy(), np_1darray_object), np_1darray_object)
 
     s_bytes = pd.Series(["a", "b", "c"]).astype(bytes)
     check(assert_type(s_bytes.to_numpy(), np_1darray_bytes), np_1darray, np.bytes_)