Skip to content

Commit b7163c2

Browse files
ramvikramsramvikrams
authored
added np.timedelta64 for series arithmatic methods (#432)
* update * update * Update test_timefuncs.py * Update test_timefuncs.py * Update test_timefuncs.py * Update test_timefuncs.py * Update test_timefuncs.py * Update test_timefuncs.py * Update test_timefuncs.py * Update test_timefuncs.py * update * update * update * update * Update series.pyi * Update series.pyi * Update test_timefuncs.py * Update test_timefuncs.py * update * update * Update test_timefuncs.py
1 parent 44f804c commit b7163c2

File tree

2 files changed

+51
-11
lines changed

2 files changed

+51
-11
lines changed

‎pandas-stubs/core/series.pyi

+30-9
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,10 @@ from pandas.core.window.rolling import (
6161
Rolling,
6262
Window,
6363
)
64-
from typing_extensions import TypeAlias
64+
from typing_extensions import (
65+
Never,
66+
TypeAlias,
67+
)
6568
import xarray as xr
6669

6770
from pandas._libs.missing import NAType
@@ -1211,7 +1214,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
12111214
def __add__(self, other: Timestamp) -> TimestampSeries: ...
12121215
@overload
12131216
def __add__(
1214-
self, other: num | _str | Timedelta | _ListLike | Series[S1]
1217+
self, other: num | _str | Timedelta | _ListLike | Series[S1] | np.timedelta64
12151218
) -> Series: ...
12161219
# ignore needed for mypy as we want different results based on the arguments
12171220
@overload
@@ -1243,7 +1246,9 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
12431246
def __le__(self, other: S1 | _ListLike | Series[S1]) -> Series[_bool]: ...
12441247
def __lt__(self, other: S1 | _ListLike | Series[S1]) -> Series[_bool]: ...
12451248
@overload
1246-
def __mul__(self, other: Timedelta | TimedeltaSeries) -> TimedeltaSeries: ...
1249+
def __mul__(
1250+
self, other: Timedelta | TimedeltaSeries | np.timedelta64
1251+
) -> TimedeltaSeries: ...
12471252
@overload
12481253
def __mul__(self, other: num | _ListLike | Series) -> Series: ...
12491254
def __mod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
@@ -1304,17 +1309,19 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
13041309
) -> TimedeltaSeries: ...
13051310
@overload
13061311
def __sub__(
1307-
self: Series[Timestamp], other: Timedelta | TimedeltaSeries | TimedeltaIndex
1312+
self: Series[Timestamp],
1313+
other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64,
13081314
) -> TimestampSeries: ...
13091315
@overload
13101316
def __sub__(
1311-
self: Series[Timedelta], other: Timedelta | TimedeltaSeries | TimedeltaIndex
1317+
self: Series[Timedelta],
1318+
other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64,
13121319
) -> TimedeltaSeries: ...
13131320
@overload
13141321
def __sub__(self, other: num | _ListLike | Series) -> Series: ...
13151322
@overload
13161323
def __truediv__(
1317-
self, other: Timedelta | TimedeltaSeries | TimedeltaIndex
1324+
self, other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64
13181325
) -> Series[float]: ...
13191326
@overload
13201327
def __truediv__(self, other: num | _ListLike | Series[S1]) -> Series: ...
@@ -1739,6 +1746,14 @@ class TimestampSeries(Series[Timestamp]):
17391746
# ignore needed because of mypy
17401747
@property
17411748
def dt(self) -> TimestampProperties: ... # type: ignore[override]
1749+
@overload # type: ignore[override]
1750+
def __add__(
1751+
self, other: TimedeltaSeries | np.timedelta64 | TimestampSeries
1752+
) -> TimestampSeries: ...
1753+
@overload
1754+
def __add__(self, other: Timestamp) -> Never: ...
1755+
def __mul__(self, other: TimestampSeries | np.timedelta64 | TimedeltaSeries) -> Never: ... # type: ignore[override]
1756+
def __truediv__(self, other: TimestampSeries | np.timedelta64 | TimedeltaSeries) -> Never: ... # type: ignore[override]
17421757

17431758
class TimedeltaSeries(Series[Timedelta]):
17441759
# ignores needed because of mypy
@@ -1747,12 +1762,18 @@ class TimedeltaSeries(Series[Timedelta]):
17471762
@overload
17481763
def __add__(self, other: Timestamp | DatetimeIndex) -> TimestampSeries: ...
17491764
@overload
1750-
def __add__(self, other: Timedelta) -> TimedeltaSeries: ...
1765+
def __add__(self, other: Timedelta | np.timedelta64) -> TimedeltaSeries: ...
17511766
def __radd__(self, pther: Timestamp | TimestampSeries) -> TimestampSeries: ... # type: ignore[override]
1752-
def __mul__(self, other: num) -> TimedeltaSeries: ... # type: ignore[override]
1767+
@overload # type: ignore[override]
1768+
def __mul__(
1769+
self, other: TimestampSeries | np.timedelta64 | Timedelta | TimedeltaSeries
1770+
) -> Never: ...
1771+
@overload
1772+
def __mul__(self, other: num) -> TimedeltaSeries: ...
17531773
def __sub__( # type: ignore[override]
1754-
self, other: Timedelta | TimedeltaSeries | TimedeltaIndex
1774+
self, other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64
17551775
) -> TimedeltaSeries: ...
1776+
def __truediv__(self, other: TimedeltaSeries | np.timedelta64 | TimedeltaIndex) -> Series[float]: ... # type: ignore[override]
17561777
@property
17571778
def dt(self) -> TimedeltaProperties: ... # type: ignore[override]
17581779

‎tests/test_timefuncs.py

+21-2
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@
1414
import pandas as pd
1515
from pandas.core.indexes.numeric import IntegerIndex
1616
import pytz
17-
from typing_extensions import assert_type
17+
from typing_extensions import (
18+
Never,
19+
assert_type,
20+
)
1821

1922
from pandas._libs import NaTType
2023
from pandas._libs.tslibs import BaseOffset
@@ -260,7 +263,7 @@ def fail_on_adding_two_timestamps() -> None:
260263
if TYPE_CHECKING_INVALID_USAGE:
261264
ssum: pd.Series = s1 + s2 # TODO both: ignore[operator]
262265
ts = pd.Timestamp("2022-06-30")
263-
tsum: pd.Series = s1 + ts # TODO both: ignore[operator]
266+
tsum: pd.Series = s1 + ts # pyright: ignore
264267

265268

266269
def test_dtindex_tzinfo() -> None:
@@ -1035,3 +1038,19 @@ def test_timedelta_range() -> None:
10351038
def test_dateoffset_freqstr() -> None:
10361039
offset = DateOffset(minutes=10)
10371040
check(assert_type(offset.freqstr, str), str)
1041+
1042+
1043+
def test_timedelta64_and_arithmatic_operator() -> None:
1044+
s1 = pd.Series(data=pd.date_range("1/1/2020", "2/1/2020"))
1045+
s2 = pd.Series(data=pd.date_range("1/1/2021", "2/1/2021"))
1046+
s3 = s2 - s1
1047+
td = np.timedelta64(1, "M")
1048+
check(assert_type((s1 - td), "TimestampSeries"), pd.Series, pd.Timestamp)
1049+
check(assert_type((s1 + td), "TimestampSeries"), pd.Series, pd.Timestamp)
1050+
check(assert_type((s3 - td), "TimedeltaSeries"), pd.Series, pd.Timedelta)
1051+
check(assert_type((s3 + td), "TimedeltaSeries"), pd.Series, pd.Timedelta)
1052+
check(assert_type((s3 / td), "pd.Series[float]"), pd.Series, float)
1053+
if TYPE_CHECKING_INVALID_USAGE:
1054+
assert_type((s1 * td), Never) # pyright: ignore
1055+
assert_type((s1 / td), Never) # pyright: ignore
1056+
assert_type((s3 * td), Never) # pyright: ignore

0 commit comments

Comments
 (0)