Skip to content

Add numpy array shapes for return types #1325

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Aug 17, 2025
Merged

Add numpy array shapes for return types #1325

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
7bdfb49
Add numpy array shapes for return types
hamdanal Aug 15, 2025
ac4d313
Fix errors on py310
hamdanal Aug 15, 2025
6be36cf
Try separate check and assert_type
hamdanal Aug 15, 2025
c31b35c
Revert "Try separate check and assert_type"
hamdanal Aug 15, 2025
39c9fff
Skip mypy failing tests on Python 3.10
hamdanal Aug 15, 2025
2d4ccf1
Add comment
hamdanal Aug 15, 2025
abfef56
See if renaming variables work
hamdanal Aug 16, 2025
519be94
Type timedelta array shapes and refactor ts and td tests
hamdanal Aug 16, 2025
252c7e9
Period also supports arrays
hamdanal Aug 16, 2025
ff5e013
Fix pyright on Python 3.10
hamdanal Aug 16, 2025
8f43a61
Add dtype to index and series subclasses
hamdanal Aug 17, 2025
fd5f866
Ignore pyrefly override
hamdanal Aug 17, 2025
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
26 changes: 19 additions & 7 deletions pandas-stubs/_libs/interval.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ from pandas.core.series import (
from pandas._typing import (
IntervalClosedType,
IntervalT,
np_ndarray_bool,
np_1darray,
npt,
)

Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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
Expand Down
38 changes: 31 additions & 7 deletions pandas-stubs/_libs/tslibs/period.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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): ...

Expand Down Expand Up @@ -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]
Expand Down
95 changes: 55 additions & 40 deletions pandas-stubs/_libs/tslibs/timedeltas.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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: ...
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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: ...
Expand Down
Loading
Loading