Add numpy array shapes for return types

pandas-stubs/_libs/tslibs/period.pyi

@@ -106,7 +106,7 @@ class Period(PeriodMixin):
     @overload
     def __eq__(self, other: PeriodSeries) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
-    def __eq__(self, other: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ...  # type: ignore[overload-overlap]
+    def __eq__(self, other: np_ndarray[ShapeT, np.object_]) -> np_ndarray[ShapeT, np.bool]: ...  # type: ignore[overload-overlap]
     @overload
     def __eq__(self, other: object) -> Literal[False]: ...
     @overload
@@ -116,31 +116,39 @@ class Period(PeriodMixin):
     @overload
     def __ge__(self, other: PeriodSeries) -> Series[bool]: ...
     @overload
-    def __ge__(self, other: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ...
+    def __ge__(
+        self, other: np_ndarray[ShapeT, np.object_]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
     @overload
     def __gt__(self, other: Period) -> bool: ...
     @overload
     def __gt__(self, other: PeriodIndex) -> np_1darray[np.bool]: ...
     @overload
     def __gt__(self, other: PeriodSeries) -> Series[bool]: ...
     @overload
-    def __gt__(self, other: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ...
+    def __gt__(
+        self, other: np_ndarray[ShapeT, np.object_]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
     @overload
     def __le__(self, other: Period) -> bool: ...
     @overload
     def __le__(self, other: PeriodIndex) -> np_1darray[np.bool]: ...
     @overload
     def __le__(self, other: PeriodSeries) -> Series[bool]: ...
     @overload
-    def __le__(self, other: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ...
+    def __le__(
+        self, other: np_ndarray[ShapeT, np.object_]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
     @overload
     def __lt__(self, other: Period) -> bool: ...
     @overload
     def __lt__(self, other: PeriodIndex) -> np_1darray[np.bool]: ...
     @overload
     def __lt__(self, other: PeriodSeries) -> Series[bool]: ...
     @overload
-    def __lt__(self, other: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ...
+    def __lt__(
+        self, other: np_ndarray[ShapeT, np.object_]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
     #  ignore[misc] here because we know all other comparisons
     #  are False, so we use Literal[False]
     @overload
@@ -150,7 +158,7 @@ class Period(PeriodMixin):
     @overload
     def __ne__(self, other: PeriodSeries) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
-    def __ne__(self, other: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ...  # type: ignore[overload-overlap]
+    def __ne__(self, other: np_ndarray[ShapeT, np.object_]) -> np_ndarray[ShapeT, np.bool]: ...  # type: ignore[overload-overlap]
     @overload
     def __ne__(self, other: object) -> Literal[True]: ...
     # Ignored due to indecipherable error from mypy:

pandas-stubs/_typing.pyi

@@ -821,13 +821,21 @@ np_ndarray_str: TypeAlias = npt.NDArray[np.str_]
 
 # Define shape and generic type variables with defaults similar to numpy
 GenericT = TypeVar("GenericT", bound=np.generic, default=Any)
+GenericT_co = TypeVar("GenericT_co", bound=np.generic, default=Any, covariant=True)
 ShapeT = TypeVar("ShapeT", bound=tuple[int, ...], default=tuple[Any, ...])
 # Numpy ndarray with more ergonomic typevar
 np_ndarray: TypeAlias = np.ndarray[ShapeT, np.dtype[GenericT]]
 # Numpy arrays with known shape (Do not use as argument types, only as return types)
 np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[GenericT]]
 np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]]
 
+class SupportsDType(Protocol[GenericT_co]):
+    @property
+    def dtype(self) -> np.dtype[GenericT_co]: ...
+
+# Similar to npt.DTypeLike but leaves out np.dtype and None for use in overloads
+DTypeLike: TypeAlias = type[Any] | tuple[Any, Any] | list[Any] | str
+
 IndexType: TypeAlias = slice | np_ndarray_anyint | Index | list[int] | Series[int]
 MaskType: TypeAlias = Series[bool] | np_ndarray_bool | list[bool]
 

pandas-stubs/core/base.pyi

@@ -24,10 +24,13 @@ from pandas._typing import (
     S1,
     AxisIndex,
     DropKeep,
+    DTypeLike,
+    GenericT,
+    GenericT_co,
     NDFrameT,
     Scalar,
+    SupportsDType,
     np_1darray,
-    npt,
 )
 from pandas.util._decorators import cache_readonly
 
@@ -43,7 +46,7 @@ class SelectionMixin(Generic[NDFrameT]):
     def __getitem__(self, key): ...
     def aggregate(self, func, *args, **kwargs): ...
 
-class IndexOpsMixin(OpsMixin, Generic[S1]):
+class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
     __array_priority__: int = ...
     @property
     def T(self) -> Self: ...
@@ -58,9 +61,26 @@ class IndexOpsMixin(OpsMixin, Generic[S1]):
     def size(self) -> int: ...
     @property
     def array(self) -> ExtensionArray: ...
+    @overload
+    def to_numpy(
+        self,
+        dtype: None = None,
+        copy: bool = False,
+        na_value: Scalar = ...,
+        **kwargs,
+    ) -> np_1darray[GenericT_co]: ...
+    @overload
+    def to_numpy(
+        self,
+        dtype: np.dtype[GenericT] | SupportsDType[GenericT] | type[GenericT],
+        copy: bool = False,
+        na_value: Scalar = ...,
+        **kwargs,
+    ) -> np_1darray[GenericT]: ...
+    @overload
     def to_numpy(
         self,
-        dtype: npt.DTypeLike | None = ...,
+        dtype: DTypeLike,
         copy: bool = False,
         na_value: Scalar = ...,
         **kwargs,

pandas-stubs/core/indexes/base.pyi

@@ -13,10 +13,12 @@ from datetime import (
 from typing import (
     Any,
     ClassVar,
+    Generic,
     Literal,
     TypeAlias,
     final,
     overload,
+    type_check_only,
 )
 
 import numpy as np
@@ -49,7 +51,10 @@ from pandas._typing import (
     AxesData,
     DropKeep,
     DtypeArg,
+    DTypeLike,
     DtypeObj,
+    GenericT,
+    GenericT_co,
     HashableT,
     IgnoreRaise,
     Label,
@@ -60,6 +65,7 @@ from pandas._typing import (
     Scalar,
     SequenceNotStr,
     SliceType,
+    SupportsDType,
     TimedeltaDtypeArg,
     TimestampDtypeArg,
     np_1darray,
@@ -519,3 +525,30 @@ class Index(IndexOpsMixin[S1]):
     def infer_objects(self, copy: bool = True) -> Self: ...
 
 UnknownIndex: TypeAlias = Index[Any]
+
+@type_check_only
+class _IndexSubclassBase(Index[S1], Generic[S1, GenericT_co]):
+    @overload
+    def to_numpy(  # pyrefly: ignore
+        self,
+        dtype: None = None,
+        copy: bool = False,
+        na_value: Scalar = ...,
+        **kwargs,
+    ) -> np_1darray[GenericT_co]: ...
+    @overload
+    def to_numpy(
+        self,
+        dtype: np.dtype[GenericT] | SupportsDType[GenericT] | type[GenericT],
+        copy: bool = False,
+        na_value: Scalar = ...,
+        **kwargs,
+    ) -> np_1darray[GenericT]: ...
+    @overload
+    def to_numpy(
+        self,
+        dtype: DTypeLike,
+        copy: bool = False,
+        na_value: Scalar = ...,
+        **kwargs,
+    ) -> np_1darray: ...

pandas-stubs/core/indexes/datetimelike.pyi

@@ -7,10 +7,11 @@ from pandas._libs.tslibs import BaseOffset
 from pandas._typing import (
     S1,
     AxisIndex,
+    GenericT_co,
     TimeUnit,
 )
 
-class DatetimeIndexOpsMixin(ExtensionIndex[S1]):
+class DatetimeIndexOpsMixin(ExtensionIndex[S1, GenericT_co]):
     @property
     def freq(self) -> BaseOffset | None: ...
     @property
@@ -33,7 +34,7 @@ class DatetimeIndexOpsMixin(ExtensionIndex[S1]):
         self, other: DatetimeIndexOpsMixin
     ) -> TimedeltaIndex: ...
 
-class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]):
+class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1, GenericT_co]):
     @property
     def unit(self) -> TimeUnit: ...
     def as_unit(self, unit: TimeUnit) -> Self: ...

pandas-stubs/core/indexes/datetimes.pyi

@@ -43,7 +43,9 @@ from pandas.core.dtypes.dtypes import DatetimeTZDtype
 
 from pandas.tseries.offsets import BaseOffset
 
-class DatetimeIndex(DatetimeTimedeltaMixin[Timestamp], DatetimeIndexProperties):
+class DatetimeIndex(
+    DatetimeTimedeltaMixin[Timestamp, np.datetime64], DatetimeIndexProperties
+):
     def __new__(
         cls,
         data: AxesData,

pandas-stubs/core/indexes/extension.pyi

@@ -1,5 +1,8 @@
-from pandas.core.indexes.base import Index
+from pandas.core.indexes.base import _IndexSubclassBase
 
-from pandas._typing import S1
+from pandas._typing import (
+    S1,
+    GenericT_co,
+)
 
-class ExtensionIndex(Index[S1]): ...
+class ExtensionIndex(_IndexSubclassBase[S1, GenericT_co]): ...

pandas-stubs/core/indexes/interval.pyi

@@ -66,7 +66,7 @@ _EdgesTimedelta: TypeAlias = (
 _TimestampLike: TypeAlias = pd.Timestamp | np.datetime64 | dt.datetime
 _TimedeltaLike: TypeAlias = pd.Timedelta | np.timedelta64 | dt.timedelta
 
-class IntervalIndex(ExtensionIndex[IntervalT], IntervalMixin):
+class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin):
     closed: IntervalClosedType
 
     def __new__(

pandas-stubs/core/indexes/period.pyi

@@ -5,6 +5,7 @@ from typing import (
     overload,
 )
 
+import numpy as np
 import pandas as pd
 from pandas import Index
 from pandas.core.indexes.accessors import PeriodIndexFieldOps
@@ -19,7 +20,7 @@ from pandas._libs.tslibs import (
 )
 from pandas._libs.tslibs.period import _PeriodAddSub
 
-class PeriodIndex(DatetimeIndexOpsMixin[pd.Period], PeriodIndexFieldOps):
+class PeriodIndex(DatetimeIndexOpsMixin[pd.Period, np.object_], PeriodIndexFieldOps):
     def __new__(
         cls,
         data=...,

pandas-stubs/core/indexes/range.pyi

@@ -8,7 +8,10 @@ from typing import (
 )
 
 import numpy as np
-from pandas.core.indexes.base import Index
+from pandas.core.indexes.base import (
+    Index,
+    _IndexSubclassBase,
+)
 
 from pandas._typing import (
     HashableT,
@@ -17,7 +20,7 @@ from pandas._typing import (
     np_ndarray_anyint,
 )
 
-class RangeIndex(Index[int]):
+class RangeIndex(_IndexSubclassBase[int, np.int64]):
     def __new__(
         cls,
         start: int | RangeIndex | range = ...,

pandas-stubs/core/indexes/timedeltas.pyi

@@ -33,7 +33,9 @@ from pandas._typing import (
     num,
 )
 
-class TimedeltaIndex(DatetimeTimedeltaMixin[Timedelta], TimedeltaIndexProperties):
+class TimedeltaIndex(
+    DatetimeTimedeltaMixin[Timedelta, np.timedelta64], TimedeltaIndexProperties
+):
     def __new__(
         cls,
         data: (