From 5dba5336ce8e727ded5932405ed454b79e35c6dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= Date: Fri, 18 Apr 2025 21:40:07 -0400 Subject: [PATCH 1/2] GH1189/GH1181 Timestamp subtraction and index shift --- pandas-stubs/_libs/tslibs/timestamps.pyi | 2 ++ pandas-stubs/core/indexes/base.pyi | 2 +- tests/test_indexes.py | 15 +++++++++++++++ tests/test_timefuncs.py | 8 ++++++++ 4 files changed, 26 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index 76416b7de..936567fdd 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/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]: ... diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 2708e2247..2aa3bb04d 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -393,7 +393,7 @@ 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 shift(self, periods: int = ..., freq=...) -> Self: ... def argsort(self, *args, **kwargs): ... def get_indexer_non_unique(self, target): ... def get_indexer_for(self, target, **kwargs): ... diff --git a/tests/test_indexes.py b/tests/test_indexes.py index bb54d11d0..f276fdba1 100644 --- a/tests/test_indexes.py +++ b/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) diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 40ee32959..ab5a3fe75 100644 --- a/tests/test_timefuncs.py +++ b/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) From 4fdd1115981b5d54480ef482ee671765ccaa21ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= Date: Sat, 19 Apr 2025 13:50:57 -0400 Subject: [PATCH 2/2] GH1189 PR feedback --- pandas-stubs/core/indexes/base.pyi | 1 - pandas-stubs/core/indexes/datetimes.pyi | 2 ++ pandas-stubs/core/indexes/period.pyi | 1 + pandas-stubs/core/indexes/timedeltas.pyi | 1 + 4 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 2aa3bb04d..2cef882ec 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/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=...) -> Self: ... def argsort(self, *args, **kwargs): ... def get_indexer_non_unique(self, target): ... def get_indexer_for(self, target, **kwargs): ... diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index 5db9e706f..c1f596c48 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/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 = ..., diff --git a/pandas-stubs/core/indexes/period.pyi b/pandas-stubs/core/indexes/period.pyi index 88cea196a..ae16b70bd 100644 --- a/pandas-stubs/core/indexes/period.pyi +++ b/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: ( diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index 93bf32aff..59b2c7aa3 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/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 = ...,