Add numpy array shapes for return types (#1325)

* Add numpy array shapes for return types

* Fix errors on py310

* Try separate check and assert_type

* Revert "Try separate check and assert_type"

This reverts commit 6be36cf.

* Skip mypy failing tests on Python 3.10

* Add comment

* See if renaming variables work

* Type timedelta array shapes and refactor ts and td tests

* Period also supports arrays

* Fix pyright on Python 3.10

* Add dtype to index and series subclasses

* Ignore pyrefly override

Author: hamdanal

pandas-stubs/_libs/interval.pyi

@@ -21,7 +21,7 @@ from pandas.core.series import (
 from pandas._typing import (
     IntervalClosedType,
     IntervalT,
-    np_ndarray_bool,
+    np_1darray,
     npt,
 )
 
@@ -170,7 +170,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __gt__(self, other: Interval[_OrderableT]) -> bool: ...
     @overload
-    def __gt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __gt__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __gt__(
         self,
@@ -179,7 +181,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __lt__(self, other: Interval[_OrderableT]) -> bool: ...
     @overload
-    def __lt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __lt__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __lt__(
         self,
@@ -188,7 +192,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __ge__(self, other: Interval[_OrderableT]) -> bool: ...
     @overload
-    def __ge__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __ge__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __ge__(
         self,
@@ -197,19 +203,25 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __le__(self, other: Interval[_OrderableT]) -> bool: ...
     @overload
-    def __le__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __le__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __eq__(self, other: Interval[_OrderableT]) -> bool: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
-    def __eq__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __eq__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __eq__(self, other: Series[_OrderableT]) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
     def __eq__(self, other: object) -> Literal[False]: ...
     @overload
     def __ne__(self, other: Interval[_OrderableT]) -> bool: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
-    def __ne__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __ne__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __ne__(self, other: Series[_OrderableT]) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload

pandas-stubs/_libs/tslibs/period.pyi

@@ -22,7 +22,11 @@ from typing_extensions import TypeAlias
 from pandas._libs.tslibs import NaTType
 from pandas._libs.tslibs.offsets import BaseOffset
 from pandas._libs.tslibs.timestamps import Timestamp
-from pandas._typing import npt
+from pandas._typing import (
+    ShapeT,
+    np_1darray,
+    np_ndarray,
+)
 
 class IncompatibleFrequency(ValueError): ...
 
@@ -98,44 +102,64 @@ class Period(PeriodMixin):
     @overload
     def __eq__(self, other: Period) -> bool: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
-    def __eq__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...  # type: ignore[overload-overlap]
+    def __eq__(self, other: PeriodIndex) -> np_1darray[np.bool]: ...  # type: ignore[overload-overlap]
     @overload
     def __eq__(self, other: PeriodSeries) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
+    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
     def __ge__(self, other: Period) -> bool: ...
     @overload
-    def __ge__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
+    def __ge__(self, other: PeriodIndex) -> np_1darray[np.bool]: ...
     @overload
     def __ge__(self, other: PeriodSeries) -> Series[bool]: ...
     @overload
+    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) -> npt.NDArray[np.bool_]: ...
+    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.object_]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
+    @overload
     def __le__(self, other: Period) -> bool: ...
     @overload
-    def __le__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
+    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.object_]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
+    @overload
     def __lt__(self, other: Period) -> bool: ...
     @overload
-    def __lt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
+    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.object_]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
     #  ignore[misc] here because we know all other comparisons
     #  are False, so we use Literal[False]
     @overload
     def __ne__(self, other: Period) -> bool: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
-    def __ne__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...  # type: ignore[overload-overlap]
+    def __ne__(self, other: PeriodIndex) -> np_1darray[np.bool]: ...  # type: ignore[overload-overlap]
     @overload
     def __ne__(self, other: PeriodSeries) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
+    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:
     # Forward operator "__add__" is not callable  [misc]

pandas-stubs/_libs/tslibs/timedeltas.pyi

@@ -33,7 +33,10 @@ from pandas._libs.tslibs import (
 from pandas._libs.tslibs.period import Period
 from pandas._libs.tslibs.timestamps import Timestamp
 from pandas._typing import (
+    ShapeT,
     TimeUnit,
+    np_1darray,
+    np_ndarray,
     npt,
 )
 
@@ -153,12 +156,12 @@ class Timedelta(timedelta):
     def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ...
     @overload
     def __add__(
-        self, other: npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.timedelta64]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.timedelta64]: ...
     @overload
     def __add__(
-        self, other: npt.NDArray[np.datetime64]
-    ) -> npt.NDArray[np.datetime64]: ...
+        self, other: np_ndarray[ShapeT, np.datetime64]
+    ) -> np_ndarray[ShapeT, np.datetime64]: ...
     @overload
     def __add__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
     @overload
@@ -176,12 +179,12 @@ class Timedelta(timedelta):
     def __radd__(self, other: NaTType) -> NaTType: ...
     @overload
     def __radd__(
-        self, other: npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.timedelta64]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.timedelta64]: ...
     @overload
     def __radd__(
-        self, other: npt.NDArray[np.datetime64]
-    ) -> npt.NDArray[np.datetime64]: ...
+        self, other: np_ndarray[ShapeT, np.datetime64]
+    ) -> np_ndarray[ShapeT, np.datetime64]: ...
     @overload
     def __radd__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
     @overload
@@ -193,8 +196,8 @@ class Timedelta(timedelta):
     def __sub__(self, other: NaTType) -> NaTType: ...
     @overload
     def __sub__(
-        self, other: npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.timedelta64]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.timedelta64]: ...
     @overload
     def __sub__(self, other: pd.TimedeltaIndex) -> TimedeltaIndex: ...
     @overload
@@ -215,12 +218,12 @@ class Timedelta(timedelta):
     def __rsub__(self, other: DatetimeIndex) -> DatetimeIndex: ...
     @overload
     def __rsub__(
-        self, other: npt.NDArray[np.datetime64]
-    ) -> npt.NDArray[np.datetime64]: ...
+        self, other: np_ndarray[ShapeT, np.datetime64]
+    ) -> np_ndarray[ShapeT, np.datetime64]: ...
     @overload
     def __rsub__(
-        self, other: npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.timedelta64]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.timedelta64]: ...
     @overload
     def __rsub__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
     def __neg__(self) -> Timedelta: ...
@@ -231,8 +234,8 @@ class Timedelta(timedelta):
     def __mul__(self, other: float) -> Timedelta: ...
     @overload
     def __mul__(
-        self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
-    ) -> npt.NDArray[np.timedelta64]: ...
+        self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
+    ) -> np_ndarray[ShapeT, np.timedelta64]: ...
     @overload
     def __mul__(self, other: Series[int]) -> TimedeltaSeries: ...
     @overload
@@ -243,8 +246,8 @@ class Timedelta(timedelta):
     def __rmul__(self, other: float) -> Timedelta: ...
     @overload
     def __rmul__(
-        self, other: npt.NDArray[np.floating] | npt.NDArray[np.integer]
-    ) -> npt.NDArray[np.timedelta64]: ...
+        self, other: np_ndarray[ShapeT, np.floating] | np_ndarray[ShapeT, np.integer]
+    ) -> np_ndarray[ShapeT, np.timedelta64]: ...
     @overload
     def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ...
     @overload
@@ -260,12 +263,12 @@ class Timedelta(timedelta):
     def __floordiv__(self, other: float) -> Timedelta: ...
     @overload
     def __floordiv__(
-        self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
-    ) -> npt.NDArray[np.timedelta64]: ...
+        self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
+    ) -> np_ndarray[ShapeT, np.timedelta64]: ...
     @overload
     def __floordiv__(
-        self, other: npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.int_]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.int_]: ...
     @overload
     def __floordiv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
     @overload
@@ -282,17 +285,17 @@ class Timedelta(timedelta):
     def __rfloordiv__(self, other: NaTType | None) -> float: ...
     @overload
     def __rfloordiv__(
-        self, other: npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.int_]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.int_]: ...
     # Override due to more types supported than dt.timedelta
     @overload  # type: ignore[override]
     def __truediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
     @overload
     def __truediv__(self, other: float) -> Timedelta: ...
     @overload
     def __truediv__(
-        self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
-    ) -> npt.NDArray[np.timedelta64]: ...
+        self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
+    ) -> np_ndarray[ShapeT, np.timedelta64]: ...
     @overload
     def __truediv__(self, other: TimedeltaSeries) -> Series[float]: ...
     @overload
@@ -308,9 +311,11 @@ class Timedelta(timedelta):
     @overload
     def __eq__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
+    def __eq__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...  # type: ignore[overload-overlap]
+    @overload
     def __eq__(  # type: ignore[overload-overlap]
-        self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.bool_]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.bool_]: ...
     @overload
     def __eq__(self, other: object) -> Literal[False]: ...
     # Override due to more types supported than dt.timedelta
@@ -319,9 +324,11 @@ class Timedelta(timedelta):
     @overload
     def __ne__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
+    def __ne__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...  # type: ignore[overload-overlap]
+    @overload
     def __ne__(  # type: ignore[overload-overlap]
-        self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.bool_]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.bool_]: ...
     @overload
     def __ne__(self, other: object) -> Literal[True]: ...
     # Override due to more types supported than dt.timedelta
@@ -335,8 +342,8 @@ class Timedelta(timedelta):
     def __mod__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
     @overload
     def __mod__(
-        self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
-    ) -> npt.NDArray[np.timedelta64]: ...
+        self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
+    ) -> np_ndarray[ShapeT, np.timedelta64]: ...
     @overload
     def __mod__(
         self, other: Series[int] | Series[float] | TimedeltaSeries
@@ -348,36 +355,44 @@ class Timedelta(timedelta):
     @overload  # type: ignore[override]
     def __le__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ...  # type: ignore[misc]
     @overload
+    def __le__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
+    @overload
     def __le__(
-        self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.bool_]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.bool_]: ...
     @overload
     def __le__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
     # Override due to more types supported than dt.timedelta
     @overload  # type: ignore[override]
     def __lt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ...  # type: ignore[misc]
     @overload
+    def __lt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
+    @overload
     def __lt__(
-        self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.bool_]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.bool_]: ...
     @overload
     def __lt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
     # Override due to more types supported than dt.timedelta
     @overload  # type: ignore[override]
     def __ge__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ...  # type: ignore[misc]
     @overload
+    def __ge__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
+    @overload
     def __ge__(
-        self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.bool_]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.bool_]: ...
     @overload
     def __ge__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
     # Override due to more types supported than dt.timedelta
     @overload  # type: ignore[override]
     def __gt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ...  # type: ignore[misc]
     @overload
+    def __gt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
+    @overload
     def __gt__(
-        self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.bool_]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.bool_]: ...
     @overload
     def __gt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
     def __hash__(self) -> int: ...