GH1189/GH1181 Timestamp subtraction and index shift (#1191)

* GH1189/GH1181 Timestamp subtraction and index shift

* GH1189 PR feedback

Author: loicdiridollou

pandas-stubs/_libs/tslibs/timestamps.pyi

@@ -227,6 +227,8 @@ class Timestamp(datetime, SupportsIndex):
     @overload
     def __sub__(self, other: TimedeltaSeries) -> TimestampSeries: ...
     @overload
+    def __sub__(self, other: TimestampSeries) -> TimedeltaSeries: ...
+    @overload
     def __sub__(
         self, other: npt.NDArray[np.timedelta64]
     ) -> npt.NDArray[np.datetime64]: ...

pandas-stubs/core/indexes/base.pyi

@@ -393,7 +393,6 @@ class Index(IndexOpsMixin[S1]):
     def asof_locs(self, where, mask): ...
     def sort_values(self, return_indexer: bool = ..., ascending: bool = ...): ...
     def sort(self, *args, **kwargs) -> None: ...
-    def shift(self, periods: int = ..., freq=...) -> None: ...
     def argsort(self, *args, **kwargs): ...
     def get_indexer_non_unique(self, target): ...
     def get_indexer_for(self, target, **kwargs): ...

pandas-stubs/core/indexes/datetimes.pyi

@@ -27,6 +27,7 @@ from pandas.core.series import (
     TimedeltaSeries,
     TimestampSeries,
 )
+from typing_extensions import Self
 
 from pandas._typing import (
     AnyArrayLike,
@@ -90,6 +91,7 @@ class DatetimeIndex(DatetimeTimedeltaMixin[Timestamp], DatetimeIndexProperties):
     def tzinfo(self) -> tzinfo | None: ...
     @property
     def dtype(self) -> np.dtype | DatetimeTZDtype: ...
+    def shift(self, periods: int = ..., freq=...) -> Self: ...
 
 def date_range(
     start: str | DateAndDatetimeLike | None = ...,

pandas-stubs/core/indexes/period.pyi

@@ -75,6 +75,7 @@ class PeriodIndex(DatetimeIndexOpsMixin[pd.Period], PeriodIndexFieldOps):
     def memory_usage(self, deep: bool = ...): ...
     @property
     def freqstr(self) -> str: ...
+    def shift(self, periods: int = ..., freq=...) -> Self: ...
 
 def period_range(
     start: (

pandas-stubs/core/indexes/timedeltas.pyi

@@ -77,6 +77,7 @@ class TimedeltaIndex(DatetimeTimedeltaMixin[Timedelta], TimedeltaIndexProperties
     def inferred_type(self) -> str: ...
     def insert(self, loc, item): ...
     def to_series(self, index=..., name: Hashable = ...) -> TimedeltaSeries: ...
+    def shift(self, periods: int = ..., freq=...) -> Self: ...
 
 def timedelta_range(
     start: TimedeltaConvertibleTypes = ...,

tests/test_indexes.py

@@ -1272,3 +1272,18 @@ def test_datetime_index_max_min_reductions() -> None:
     check(assert_type(dtidx.argmin(), np.int64), np.int64)
     check(assert_type(dtidx.max(), pd.Timestamp), pd.Timestamp)
     check(assert_type(dtidx.min(), pd.Timestamp), pd.Timestamp)
+
+
+def test_periodindex_shift() -> None:
+    ind = pd.period_range(start="2022-06-01", periods=10)
+    check(assert_type(ind.shift(1), pd.PeriodIndex), pd.PeriodIndex)
+
+
+def test_datetimeindex_shift() -> None:
+    ind = pd.date_range("2023-01-01", "2023-02-01")
+    check(assert_type(ind.shift(1), pd.DatetimeIndex), pd.DatetimeIndex)
+
+
+def test_timedeltaindex_shift() -> None:
+    ind = pd.date_range("1/1/2021", "1/5/2021") - pd.Timestamp("1/3/2019")
+    check(assert_type(ind.shift(1), pd.TimedeltaIndex), pd.TimedeltaIndex)

tests/test_timefuncs.py

@@ -1474,3 +1474,11 @@ def test_DatetimeIndex_sub_timedelta() -> None:
 def test_to_offset() -> None:
     check(assert_type(to_offset(None), None), type(None))
     check(assert_type(to_offset("1D"), DateOffset), DateOffset)
+
+
+def test_timestamp_sub_series() -> None:
+    """Test subtracting Series[Timestamp] from Timestamp (see GH1189)."""
+    ts1 = pd.to_datetime(pd.Series(["2022-03-05", "2022-03-06"]))
+    one_ts = ts1.iloc[0]
+    check(assert_type(ts1.iloc[0], pd.Timestamp), pd.Timestamp)
+    check(assert_type(one_ts - ts1, "TimedeltaSeries"), pd.Series, pd.Timedelta)