From d8abb362b071d92692f08816290a13a969fca80b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81d=C3=A1m=20Lippai?= Date: Fri, 13 Sep 2024 13:15:18 -0400 Subject: [PATCH 1/4] Add datetime.date and datetime.datetime specializations --- pandas-stubs/_libs/tslibs/offsets.pyi | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/pandas-stubs/_libs/tslibs/offsets.pyi b/pandas-stubs/_libs/tslibs/offsets.pyi index 1acdadf78..a2488b657 100644 --- a/pandas-stubs/_libs/tslibs/offsets.pyi +++ b/pandas-stubs/_libs/tslibs/offsets.pyi @@ -14,6 +14,7 @@ from typing import ( from dateutil.relativedelta import weekday as WeekdayClass import numpy as np +from pandas import Timestamp from pandas.core.indexes.datetimes import DatetimeIndex from typing_extensions import Self @@ -42,6 +43,10 @@ class BaseOffset: @overload def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload + def __add__(self, other: date) -> Timestamp: ... + @overload + def __add__(self, other: datetime) -> Timestamp: ... + @overload def __add__(self, other: BaseOffset) -> Self: ... @overload def __add__(self, other: _DatetimeT) -> _DatetimeT: ... @@ -50,15 +55,28 @@ class BaseOffset: @overload def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload + def __radd__(self, other: date) -> Timestamp: ... + @overload + def __radd__(self, other: datetime) -> Timestamp: ... + @overload def __radd__(self, other: BaseOffset) -> Self: ... @overload def __radd__(self, other: _DatetimeT) -> _DatetimeT: ... @overload def __radd__(self, other: _TimedeltaT) -> _TimedeltaT: ... + @overload + def __sub__(self, other: date) -> Timestamp: ... + @overload + def __sub__(self, other: datetime) -> Timestamp: ... + @overload def __sub__(self, other: BaseOffset) -> Self: ... @overload def __rsub__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload + def __rsub__(self, other: date) -> Timestamp: ... + @overload + def __rsub__(self, other: datetime) -> Timestamp: ... + @overload def __rsub__(self, other: BaseOffset) -> Self: ... @overload def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ... From 7a4bee7f9342128ff9c3c502fb38e237d3f3e496 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81d=C3=A1m=20Lippai?= Date: Sat, 14 Sep 2024 06:00:25 +0200 Subject: [PATCH 2/4] Add tests and update overload order --- pandas-stubs/_libs/tslibs/offsets.pyi | 13 ------------- tests/test_timefuncs.py | 6 ++++-- 2 files changed, 4 insertions(+), 15 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/offsets.pyi b/pandas-stubs/_libs/tslibs/offsets.pyi index a2488b657..f38adf879 100644 --- a/pandas-stubs/_libs/tslibs/offsets.pyi +++ b/pandas-stubs/_libs/tslibs/offsets.pyi @@ -45,8 +45,6 @@ class BaseOffset: @overload def __add__(self, other: date) -> Timestamp: ... @overload - def __add__(self, other: datetime) -> Timestamp: ... - @overload def __add__(self, other: BaseOffset) -> Self: ... @overload def __add__(self, other: _DatetimeT) -> _DatetimeT: ... @@ -57,26 +55,15 @@ class BaseOffset: @overload def __radd__(self, other: date) -> Timestamp: ... @overload - def __radd__(self, other: datetime) -> Timestamp: ... - @overload def __radd__(self, other: BaseOffset) -> Self: ... @overload def __radd__(self, other: _DatetimeT) -> _DatetimeT: ... @overload def __radd__(self, other: _TimedeltaT) -> _TimedeltaT: ... - @overload - def __sub__(self, other: date) -> Timestamp: ... - @overload - def __sub__(self, other: datetime) -> Timestamp: ... - @overload - def __sub__(self, other: BaseOffset) -> Self: ... - @overload def __rsub__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload def __rsub__(self, other: date) -> Timestamp: ... @overload - def __rsub__(self, other: datetime) -> Timestamp: ... - @overload def __rsub__(self, other: BaseOffset) -> Self: ... @overload def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ... diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index ed6cf910e..d1da094dd 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -664,8 +664,10 @@ def test_some_offsets() -> None: ), pd.DatetimeIndex, ) - # GH 224 - check(assert_type(dt.date.today() - Day(), dt.date), dt.date) + # GH 755 + check(assert_type(dt.date.today() - Day(), pd.Timestamp), pd.Timestamp) + check(assert_type(dt.date.today() + Day(), pd.Timestamp), pd.Timestamp) + check(assert_type(Day() + dt.date.today(), pd.Timestamp), pd.Timestamp) # GH 235 check( assert_type( From 1bad029ab6d1b7a215e104b48fd9196e9aaa2d4f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81d=C3=A1m=20Lippai?= Date: Sat, 14 Sep 2024 06:02:12 +0200 Subject: [PATCH 3/4] Don't remove the correct type --- pandas-stubs/_libs/tslibs/offsets.pyi | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/offsets.pyi b/pandas-stubs/_libs/tslibs/offsets.pyi index f38adf879..19bfa36c4 100644 --- a/pandas-stubs/_libs/tslibs/offsets.pyi +++ b/pandas-stubs/_libs/tslibs/offsets.pyi @@ -23,7 +23,7 @@ from pandas._typing import npt from pandas.tseries.holiday import AbstractHolidayCalendar -_DatetimeT = TypeVar("_DatetimeT", bound=date) +_DatetimeT = TypeVar("_DatetimeT", bound=datetime) _TimedeltaT = TypeVar("_TimedeltaT", bound=timedelta) prefix_mapping: dict[str, type] @@ -43,31 +43,33 @@ class BaseOffset: @overload def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload + def __add__(self, other: _DatetimeT) -> _DatetimeT: ... + @overload def __add__(self, other: date) -> Timestamp: ... @overload def __add__(self, other: BaseOffset) -> Self: ... @overload - def __add__(self, other: _DatetimeT) -> _DatetimeT: ... - @overload def __add__(self, other: _TimedeltaT) -> _TimedeltaT: ... @overload def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload + def __radd__(self, other: _DatetimeT) -> _DatetimeT: ... + @overload def __radd__(self, other: date) -> Timestamp: ... @overload def __radd__(self, other: BaseOffset) -> Self: ... @overload - def __radd__(self, other: _DatetimeT) -> _DatetimeT: ... - @overload def __radd__(self, other: _TimedeltaT) -> _TimedeltaT: ... + def __sub__(self, other: BaseOffset) -> Self: ... + @overload def __rsub__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload + def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ... + @overload def __rsub__(self, other: date) -> Timestamp: ... @overload def __rsub__(self, other: BaseOffset) -> Self: ... @overload - def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ... - @overload def __rsub__(self, other: _TimedeltaT) -> _TimedeltaT: ... def __call__(self, other): ... @overload From b23d1aa693307e3b3ba73e6ab0b534bf443875a9 Mon Sep 17 00:00:00 2001 From: Adam Lippai Date: Sun, 13 Oct 2024 19:11:17 -0400 Subject: [PATCH 4/4] Ignore overload errors https://github.com/microsoft/pylance-release/issues/6512 --- pandas-stubs/_libs/tslibs/offsets.pyi | 6 +++--- tests/test_timefuncs.py | 3 +++ 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/offsets.pyi b/pandas-stubs/_libs/tslibs/offsets.pyi index 19bfa36c4..4b6a61230 100644 --- a/pandas-stubs/_libs/tslibs/offsets.pyi +++ b/pandas-stubs/_libs/tslibs/offsets.pyi @@ -43,7 +43,7 @@ class BaseOffset: @overload def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload - def __add__(self, other: _DatetimeT) -> _DatetimeT: ... + def __add__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload def __add__(self, other: date) -> Timestamp: ... @overload @@ -53,7 +53,7 @@ class BaseOffset: @overload def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload - def __radd__(self, other: _DatetimeT) -> _DatetimeT: ... + def __radd__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload def __radd__(self, other: date) -> Timestamp: ... @overload @@ -64,7 +64,7 @@ class BaseOffset: @overload def __rsub__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... @overload - def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ... + def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload def __rsub__(self, other: date) -> Timestamp: ... @overload diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index d1da094dd..9722c314d 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -668,6 +668,9 @@ def test_some_offsets() -> None: check(assert_type(dt.date.today() - Day(), pd.Timestamp), pd.Timestamp) check(assert_type(dt.date.today() + Day(), pd.Timestamp), pd.Timestamp) check(assert_type(Day() + dt.date.today(), pd.Timestamp), pd.Timestamp) + check(assert_type(dt.datetime.now() - Day(), dt.datetime), dt.datetime) + check(assert_type(dt.datetime.now() + Day(), dt.datetime), dt.datetime) + check(assert_type(Day() + dt.datetime.now(), dt.datetime), dt.datetime) # GH 235 check( assert_type(