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: ...