diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index b017b6840..91923f9eb 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -61,7 +61,10 @@ from pandas.core.window.rolling import ( Rolling, Window, ) -from typing_extensions import TypeAlias +from typing_extensions import ( + Never, + TypeAlias, +) import xarray as xr from pandas._libs.missing import NAType @@ -1211,7 +1214,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def __add__(self, other: Timestamp) -> TimestampSeries: ... @overload def __add__( - self, other: num | _str | Timedelta | _ListLike | Series[S1] + self, other: num | _str | Timedelta | _ListLike | Series[S1] | np.timedelta64 ) -> Series: ... # ignore needed for mypy as we want different results based on the arguments @overload @@ -1243,7 +1246,9 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def __le__(self, other: S1 | _ListLike | Series[S1]) -> Series[_bool]: ... def __lt__(self, other: S1 | _ListLike | Series[S1]) -> Series[_bool]: ... @overload - def __mul__(self, other: Timedelta | TimedeltaSeries) -> TimedeltaSeries: ... + def __mul__( + self, other: Timedelta | TimedeltaSeries | np.timedelta64 + ) -> TimedeltaSeries: ... @overload def __mul__(self, other: num | _ListLike | Series) -> Series: ... def __mod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... @@ -1304,17 +1309,19 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): ) -> TimedeltaSeries: ... @overload def __sub__( - self: Series[Timestamp], other: Timedelta | TimedeltaSeries | TimedeltaIndex + self: Series[Timestamp], + other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64, ) -> TimestampSeries: ... @overload def __sub__( - self: Series[Timedelta], other: Timedelta | TimedeltaSeries | TimedeltaIndex + self: Series[Timedelta], + other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64, ) -> TimedeltaSeries: ... @overload def __sub__(self, other: num | _ListLike | Series) -> Series: ... @overload def __truediv__( - self, other: Timedelta | TimedeltaSeries | TimedeltaIndex + self, other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64 ) -> Series[float]: ... @overload def __truediv__(self, other: num | _ListLike | Series[S1]) -> Series: ... @@ -1739,6 +1746,14 @@ class TimestampSeries(Series[Timestamp]): # ignore needed because of mypy @property def dt(self) -> TimestampProperties: ... # type: ignore[override] + @overload # type: ignore[override] + def __add__( + self, other: TimedeltaSeries | np.timedelta64 | TimestampSeries + ) -> TimestampSeries: ... + @overload + def __add__(self, other: Timestamp) -> Never: ... + def __mul__(self, other: TimestampSeries | np.timedelta64 | TimedeltaSeries) -> Never: ... # type: ignore[override] + def __truediv__(self, other: TimestampSeries | np.timedelta64 | TimedeltaSeries) -> Never: ... # type: ignore[override] class TimedeltaSeries(Series[Timedelta]): # ignores needed because of mypy @@ -1747,12 +1762,18 @@ class TimedeltaSeries(Series[Timedelta]): @overload def __add__(self, other: Timestamp | DatetimeIndex) -> TimestampSeries: ... @overload - def __add__(self, other: Timedelta) -> TimedeltaSeries: ... + def __add__(self, other: Timedelta | np.timedelta64) -> TimedeltaSeries: ... def __radd__(self, pther: Timestamp | TimestampSeries) -> TimestampSeries: ... # type: ignore[override] - def __mul__(self, other: num) -> TimedeltaSeries: ... # type: ignore[override] + @overload # type: ignore[override] + def __mul__( + self, other: TimestampSeries | np.timedelta64 | Timedelta | TimedeltaSeries + ) -> Never: ... + @overload + def __mul__(self, other: num) -> TimedeltaSeries: ... def __sub__( # type: ignore[override] - self, other: Timedelta | TimedeltaSeries | TimedeltaIndex + self, other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64 ) -> TimedeltaSeries: ... + def __truediv__(self, other: TimedeltaSeries | np.timedelta64 | TimedeltaIndex) -> Series[float]: ... # type: ignore[override] @property def dt(self) -> TimedeltaProperties: ... # type: ignore[override] diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 62eebf6a4..b97369466 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -14,7 +14,10 @@ import pandas as pd from pandas.core.indexes.numeric import IntegerIndex import pytz -from typing_extensions import assert_type +from typing_extensions import ( + Never, + assert_type, +) from pandas._libs import NaTType from pandas._libs.tslibs import BaseOffset @@ -260,7 +263,7 @@ def fail_on_adding_two_timestamps() -> None: if TYPE_CHECKING_INVALID_USAGE: ssum: pd.Series = s1 + s2 # TODO both: ignore[operator] ts = pd.Timestamp("2022-06-30") - tsum: pd.Series = s1 + ts # TODO both: ignore[operator] + tsum: pd.Series = s1 + ts # pyright: ignore def test_dtindex_tzinfo() -> None: @@ -1035,3 +1038,19 @@ def test_timedelta_range() -> None: def test_dateoffset_freqstr() -> None: offset = DateOffset(minutes=10) check(assert_type(offset.freqstr, str), str) + + +def test_timedelta64_and_arithmatic_operator() -> None: + s1 = pd.Series(data=pd.date_range("1/1/2020", "2/1/2020")) + s2 = pd.Series(data=pd.date_range("1/1/2021", "2/1/2021")) + s3 = s2 - s1 + td = np.timedelta64(1, "M") + check(assert_type((s1 - td), "TimestampSeries"), pd.Series, pd.Timestamp) + check(assert_type((s1 + td), "TimestampSeries"), pd.Series, pd.Timestamp) + check(assert_type((s3 - td), "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type((s3 + td), "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type((s3 / td), "pd.Series[float]"), pd.Series, float) + if TYPE_CHECKING_INVALID_USAGE: + assert_type((s1 * td), Never) # pyright: ignore + assert_type((s1 / td), Never) # pyright: ignore + assert_type((s3 * td), Never) # pyright: ignore