Skip to content

Commit f4888eb

Browse files
author
Kevin Sheppard
committed
ENH: Add NaTType for add and sub
1 parent 311a199 commit f4888eb

File tree

2 files changed

+16
-0
lines changed

2 files changed

+16
-0
lines changed

pandas-stubs/_libs/tslibs/timestamps.pyi

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ from pandas.core.series import (
2727

2828
from pandas._libs.tslibs import (
2929
BaseOffset,
30+
NaTType,
3031
Period,
3132
Tick,
3233
Timedelta,
@@ -179,12 +180,16 @@ class Timestamp(datetime):
179180
self: _DatetimeT, other: timedelta | np.timedelta64 | Tick
180181
) -> _DatetimeT: ...
181182
@overload
183+
def __add__(self, other: NaTType) -> NaTType: ...
184+
@overload
182185
def __add__(self, other: Series) -> TimestampSeries: ...
183186
@overload
184187
def __add__(self, other: TimedeltaIndex) -> DatetimeIndex: ...
185188
@overload
186189
def __radd__(self: _DatetimeT, other: timedelta) -> _DatetimeT: ...
187190
@overload
191+
def __radd__(self, other: NaTType) -> NaTType: ...
192+
@overload
188193
def __radd__(self, other: TimedeltaIndex) -> DatetimeIndex: ...
189194
@overload
190195
def __radd__(
@@ -197,13 +202,16 @@ class Timestamp(datetime):
197202
self: _DatetimeT, other: timedelta | np.timedelta64 | Tick
198203
) -> _DatetimeT: ...
199204
@overload
205+
def __sub__(self, other: NaTType) -> NaTType: ...
206+
@overload
200207
def __sub__(self, other: TimedeltaIndex) -> DatetimeIndex: ...
201208
@overload
202209
def __sub__(self, other: TimedeltaSeries) -> TimestampSeries: ...
203210
@overload
204211
def __sub__(
205212
self, other: npt.NDArray[np.timedelta64]
206213
) -> npt.NDArray[np.datetime64]: ...
214+
def __rsub__(self, other: NaTType) -> NaTType: ...
207215
@overload # type: ignore[override]
208216
def __eq__(self, other: Timestamp | np.datetime64 | datetime) -> bool: ...
209217
@overload

tests/test_scalars.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
import pytz
1616
from typing_extensions import assert_type
1717

18+
from pandas._libs.tslibs import NaTType
19+
1820
if TYPE_CHECKING:
1921
from pandas.core.series import TimedeltaSeries # noqa: F401
2022
from pandas.core.series import TimestampSeries # noqa: F401
@@ -146,13 +148,15 @@ def test_timestamp_add_sub() -> None:
146148
as4 = pd.TimedeltaIndex([1, 2, 3], "D")
147149
as5 = ts_series
148150
as6 = np_td64_arr
151+
as7 = pd.NaT
149152

150153
check(assert_type(ts + as1, pd.Timestamp), pd.Timestamp)
151154
check(assert_type(ts + as2, pd.Timestamp), pd.Timestamp)
152155
check(assert_type(ts + as3, pd.Timestamp), pd.Timestamp)
153156
check(assert_type(ts + as4, pd.DatetimeIndex), pd.DatetimeIndex)
154157
check(assert_type(ts + as5, "TimestampSeries"), pd.Series)
155158
check(assert_type(ts + as6, npt.NDArray[np.datetime64]), np.ndarray)
159+
check(assert_type(ts + as7, NaTType), NaTType)
156160

157161
check(assert_type(as1 + ts, pd.Timestamp), pd.Timestamp)
158162
check(assert_type(as2 + ts, pd.Timestamp), pd.Timestamp)
@@ -168,13 +172,17 @@ def test_timestamp_add_sub() -> None:
168172
),
169173
np.ndarray,
170174
)
175+
check(assert_type(as7 + ts, NaTType), NaTType)
171176

172177
check(assert_type(ts - as1, pd.Timestamp), pd.Timestamp)
173178
check(assert_type(ts - as2, pd.Timestamp), pd.Timestamp)
174179
check(assert_type(ts - as3, pd.Timestamp), pd.Timestamp)
175180
check(assert_type(ts - as4, pd.DatetimeIndex), pd.DatetimeIndex)
176181
check(assert_type(ts - as5, "TimestampSeries"), pd.Series)
177182
check(assert_type(ts - as6, npt.NDArray[np.datetime64]), np.ndarray)
183+
check(assert_type(ts - as7, NaTType), NaTType)
184+
185+
check(assert_type(as7 - ts, NaTType), NaTType)
178186

179187

180188
def test_timestamp_cmp() -> None:

0 commit comments

Comments
 (0)