ENH: Add NaTType for add and sub

Author: Kevin Sheppard

pandas-stubs/_libs/tslibs/timestamps.pyi

@@ -27,6 +27,7 @@ from pandas.core.series import (
 
 from pandas._libs.tslibs import (
     BaseOffset,
+    NaTType,
     Period,
     Tick,
     Timedelta,
@@ -179,12 +180,16 @@ class Timestamp(datetime):
         self: _DatetimeT, other: timedelta | np.timedelta64 | Tick
     ) -> _DatetimeT: ...
     @overload
+    def __add__(self, other: NaTType) -> NaTType: ...
+    @overload
     def __add__(self, other: Series) -> TimestampSeries: ...
     @overload
     def __add__(self, other: TimedeltaIndex) -> DatetimeIndex: ...
     @overload
     def __radd__(self: _DatetimeT, other: timedelta) -> _DatetimeT: ...
     @overload
+    def __radd__(self, other: NaTType) -> NaTType: ...
+    @overload
     def __radd__(self, other: TimedeltaIndex) -> DatetimeIndex: ...
     @overload
     def __radd__(
@@ -197,13 +202,16 @@ class Timestamp(datetime):
         self: _DatetimeT, other: timedelta | np.timedelta64 | Tick
     ) -> _DatetimeT: ...
     @overload
+    def __sub__(self, other: NaTType) -> NaTType: ...
+    @overload
     def __sub__(self, other: TimedeltaIndex) -> DatetimeIndex: ...
     @overload
     def __sub__(self, other: TimedeltaSeries) -> TimestampSeries: ...
     @overload
     def __sub__(
         self, other: npt.NDArray[np.timedelta64]
     ) -> npt.NDArray[np.datetime64]: ...
+    def __rsub__(self, other: NaTType) -> NaTType: ...
     @overload  # type: ignore[override]
     def __eq__(self, other: Timestamp | np.datetime64 | datetime) -> bool: ...
     @overload

tests/test_scalars.py

@@ -15,6 +15,8 @@
 import pytz
 from typing_extensions import assert_type
 
+from pandas._libs.tslibs import NaTType
+
 if TYPE_CHECKING:
     from pandas.core.series import TimedeltaSeries  # noqa: F401
     from pandas.core.series import TimestampSeries  # noqa: F401
@@ -146,13 +148,15 @@ def test_timestamp_add_sub() -> None:
     as4 = pd.TimedeltaIndex([1, 2, 3], "D")
     as5 = ts_series
     as6 = np_td64_arr
+    as7 = pd.NaT
 
     check(assert_type(ts + as1, pd.Timestamp), pd.Timestamp)
     check(assert_type(ts + as2, pd.Timestamp), pd.Timestamp)
     check(assert_type(ts + as3, pd.Timestamp), pd.Timestamp)
     check(assert_type(ts + as4, pd.DatetimeIndex), pd.DatetimeIndex)
     check(assert_type(ts + as5, "TimestampSeries"), pd.Series)
     check(assert_type(ts + as6, npt.NDArray[np.datetime64]), np.ndarray)
+    check(assert_type(ts + as7, NaTType), NaTType)
 
     check(assert_type(as1 + ts, pd.Timestamp), pd.Timestamp)
     check(assert_type(as2 + ts, pd.Timestamp), pd.Timestamp)
@@ -168,13 +172,17 @@ def test_timestamp_add_sub() -> None:
         ),
         np.ndarray,
     )
+    check(assert_type(as7 + ts, NaTType), NaTType)
 
     check(assert_type(ts - as1, pd.Timestamp), pd.Timestamp)
     check(assert_type(ts - as2, pd.Timestamp), pd.Timestamp)
     check(assert_type(ts - as3, pd.Timestamp), pd.Timestamp)
     check(assert_type(ts - as4, pd.DatetimeIndex), pd.DatetimeIndex)
     check(assert_type(ts - as5, "TimestampSeries"), pd.Series)
     check(assert_type(ts - as6, npt.NDArray[np.datetime64]), np.ndarray)
+    check(assert_type(ts - as7, NaTType), NaTType)
+
+    check(assert_type(as7 - ts, NaTType), NaTType)
 
 
 def test_timestamp_cmp() -> None: