Skip to content

Commit f71224c

Browse files
authored
Add datetime.date and datetime.datetime specializations (pandas-dev#1003)
* Add datetime.date and datetime.datetime specializations * Add tests and update overload order * Don't remove the correct type * Ignore overload errors microsoft/pylance-release#6512
1 parent 46eef4a commit f71224c

File tree

2 files changed

+21
-9
lines changed

2 files changed

+21
-9
lines changed

pandas-stubs/_libs/tslibs/offsets.pyi

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ from typing import (
1414

1515
from dateutil.relativedelta import weekday as WeekdayClass
1616
import numpy as np
17+
from pandas import Timestamp
1718
from pandas.core.indexes.datetimes import DatetimeIndex
1819
from typing_extensions import Self
1920

@@ -22,7 +23,7 @@ from pandas._typing import npt
2223

2324
from pandas.tseries.holiday import AbstractHolidayCalendar
2425

25-
_DatetimeT = TypeVar("_DatetimeT", bound=date)
26+
_DatetimeT = TypeVar("_DatetimeT", bound=datetime)
2627
_TimedeltaT = TypeVar("_TimedeltaT", bound=timedelta)
2728

2829
prefix_mapping: dict[str, type]
@@ -42,26 +43,32 @@ class BaseOffset:
4243
@overload
4344
def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
4445
@overload
45-
def __add__(self, other: BaseOffset) -> Self: ...
46+
def __add__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
47+
@overload
48+
def __add__(self, other: date) -> Timestamp: ...
4649
@overload
47-
def __add__(self, other: _DatetimeT) -> _DatetimeT: ...
50+
def __add__(self, other: BaseOffset) -> Self: ...
4851
@overload
4952
def __add__(self, other: _TimedeltaT) -> _TimedeltaT: ...
5053
@overload
5154
def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
5255
@overload
53-
def __radd__(self, other: BaseOffset) -> Self: ...
56+
def __radd__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
5457
@overload
55-
def __radd__(self, other: _DatetimeT) -> _DatetimeT: ...
58+
def __radd__(self, other: date) -> Timestamp: ...
59+
@overload
60+
def __radd__(self, other: BaseOffset) -> Self: ...
5661
@overload
5762
def __radd__(self, other: _TimedeltaT) -> _TimedeltaT: ...
5863
def __sub__(self, other: BaseOffset) -> Self: ...
5964
@overload
6065
def __rsub__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
6166
@overload
62-
def __rsub__(self, other: BaseOffset) -> Self: ...
67+
def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
6368
@overload
64-
def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ...
69+
def __rsub__(self, other: date) -> Timestamp: ...
70+
@overload
71+
def __rsub__(self, other: BaseOffset) -> Self: ...
6572
@overload
6673
def __rsub__(self, other: _TimedeltaT) -> _TimedeltaT: ...
6774
def __call__(self, other): ...

tests/test_timefuncs.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -664,8 +664,13 @@ def test_some_offsets() -> None:
664664
),
665665
pd.DatetimeIndex,
666666
)
667-
# GH 224
668-
check(assert_type(dt.date.today() - Day(), dt.date), dt.date)
667+
# GH 755
668+
check(assert_type(dt.date.today() - Day(), pd.Timestamp), pd.Timestamp)
669+
check(assert_type(dt.date.today() + Day(), pd.Timestamp), pd.Timestamp)
670+
check(assert_type(Day() + dt.date.today(), pd.Timestamp), pd.Timestamp)
671+
check(assert_type(dt.datetime.now() - Day(), dt.datetime), dt.datetime)
672+
check(assert_type(dt.datetime.now() + Day(), dt.datetime), dt.datetime)
673+
check(assert_type(Day() + dt.datetime.now(), dt.datetime), dt.datetime)
669674
# GH 235
670675
check(
671676
assert_type(

0 commit comments

Comments
 (0)