ENH: Improve typing for Timestamp (#389)

* ENH: Improve typing for Timestamp

* MAINT: Add dep for testing

* ENH: Improve fold

* REV: Revert Literal due to violation of Liskov subst. princip.

* TST: Add type tests for series

* CLN: Clean testing code

* ENH: Further refinements of timestamp and tests

* BUG: Correct new

* Final types and tests for timestamp

* REV: Remove mypy coverage changes

* CLN: Final clean

* CLN: Final clean

Co-authored-by: Kevin Sheppard <[email protected]>

Author: bashtage

pandas-stubs/_libs/tslibs/timestamps.pyi

@@ -8,38 +8,53 @@ from datetime import (
 from time import struct_time
 from typing import (
     ClassVar,
+    Literal,
     TypeVar,
+    Union,
     overload,
 )
 
 import numpy as np
-from pandas import Index
+from pandas import (
+    DatetimeIndex,
+    Index,
+    TimedeltaIndex,
+)
 from pandas.core.series import (
     Series,
+    TimedeltaSeries,
     TimestampSeries,
 )
+from typing_extensions import TypeAlias
 
 from pandas._libs.tslibs import (
     BaseOffset,
     Period,
     Tick,
     Timedelta,
 )
-from pandas._typing import np_ndarray_bool
+from pandas._typing import (
+    np_ndarray_bool,
+    npt,
+)
 
 _DatetimeT = TypeVar("_DatetimeT", bound=datetime)
-
-def integer_op_not_supported(obj: object) -> TypeError: ...
+_Ambiguous: TypeAlias = Union[bool, Literal["raise", "NaT"]]
+_Nonexistent: TypeAlias = Union[
+    Literal["raise", "NaT", "shift_backward", "shift_forward"], Timedelta, timedelta
+]
 
 class Timestamp(datetime):
     min: ClassVar[Timestamp]
     max: ClassVar[Timestamp]
 
     resolution: ClassVar[Timedelta]
-    value: int  # np.int64
+    value: int
     def __new__(
         cls: type[_DatetimeT],
         ts_input: np.integer | float | str | _date | datetime | np.datetime64 = ...,
+        # Freq is deprecated but is left in to allow code like Timestamp(2000,1,1)
+        # Removing it would make the other arguments position only
         freq: int | str | BaseOffset | None = ...,
         tz: str | _tzinfo | int | None = ...,
         unit: str | int | None = ...,
@@ -53,7 +68,7 @@ class Timestamp(datetime):
         nanosecond: int | None = ...,
         tzinfo: _tzinfo | None = ...,
         *,
-        fold: int | None = ...,
+        fold: Literal[0, 1] | None = ...,
     ) -> _DatetimeT: ...
     # GH 46171
     # While Timestamp can return pd.NaT, having the constructor return
@@ -73,14 +88,16 @@ class Timestamp(datetime):
     @property
     def microsecond(self) -> int: ...
     @property
+    def nanosecond(self) -> int: ...
+    @property
     def tzinfo(self) -> _tzinfo | None: ...
     @property
     def tz(self) -> _tzinfo | None: ...
     @property
     def fold(self) -> int: ...
     @classmethod
     def fromtimestamp(
-        cls: type[_DatetimeT], t: float, tz: _tzinfo | None = ...
+        cls: type[_DatetimeT], t: float, tz: _tzinfo | str | None = ...
     ) -> _DatetimeT: ...
     @classmethod
     def utcfromtimestamp(cls: type[_DatetimeT], ts: float) -> _DatetimeT: ...
@@ -90,7 +107,8 @@ class Timestamp(datetime):
     def fromordinal(
         cls: type[_DatetimeT],
         ordinal: int,
-        freq: str | BaseOffset | None = ...,
+        # freq produces a FutureWarning about being deprecated in a future version
+        freq: None = ...,
         tz: _tzinfo | str | None = ...,
     ) -> _DatetimeT: ...
     @classmethod
@@ -111,17 +129,21 @@ class Timestamp(datetime):
     def date(self) -> _date: ...
     def time(self) -> _time: ...
     def timetz(self) -> _time: ...
-    def replace(
+    # Override since fold is more precise than datetime.replace(fold:int)
+    # Here it is restricted to be 0 or 1 using a Literal
+    # Violation of Liskov substitution principle
+    def replace(  # type:ignore[override]
         self,
-        year: int = ...,
-        month: int = ...,
-        day: int = ...,
-        hour: int = ...,
-        minute: int = ...,
-        second: int = ...,
-        microsecond: int = ...,
+        year: int | None = ...,
+        month: int | None = ...,
+        day: int | None = ...,
+        hour: int | None = ...,
+        minute: int | None = ...,
+        second: int | None = ...,
+        microsecond: int | None = ...,
         tzinfo: _tzinfo | None = ...,
-        fold: int = ...,
+        *,
+        fold: Literal[0, 1] | None = ...,
     ) -> Timestamp: ...
     def astimezone(self: _DatetimeT, tz: _tzinfo | None = ...) -> _DatetimeT: ...
     def ctime(self) -> str: ...
@@ -131,45 +153,86 @@ class Timestamp(datetime):
     def utcoffset(self) -> timedelta | None: ...
     def tzname(self) -> str | None: ...
     def dst(self) -> timedelta | None: ...
+    # Mypy complains Forward operator "<inequality op>" is not callable, so ignore misc
+    # for le, lt ge and gt
     @overload  # type: ignore[override]
-    def __le__(self, other: datetime) -> bool: ...
+    def __le__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
-    def __le__(self, other: Index) -> np_ndarray_bool: ...
+    def __le__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
     @overload
-    def __le__(self, other: TimestampSeries) -> Series[bool]: ...
+    def __le__(self, other: TimestampSeries | Series[Timestamp]) -> Series[bool]: ...
     @overload  # type: ignore[override]
-    def __lt__(self, other: datetime) -> bool: ...
+    def __lt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
-    def __lt__(self, other: Index) -> np_ndarray_bool: ...
+    def __lt__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
     @overload
-    def __lt__(self, other: TimestampSeries) -> Series[bool]: ...
+    def __lt__(self, other: TimestampSeries | Series[Timestamp]) -> Series[bool]: ...
     @overload  # type: ignore[override]
-    def __ge__(self, other: datetime) -> bool: ...
+    def __ge__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
-    def __ge__(self, other: Index) -> np_ndarray_bool: ...
+    def __ge__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
     @overload
-    def __ge__(self, other: TimestampSeries) -> Series[bool]: ...
+    def __ge__(self, other: TimestampSeries | Series[Timestamp]) -> Series[bool]: ...
     @overload  # type: ignore[override]
-    def __gt__(self, other: datetime) -> bool: ...
+    def __gt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
-    def __gt__(self, other: Index) -> np_ndarray_bool: ...
+    def __gt__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
     @overload
-    def __gt__(self, other: TimestampSeries) -> Series[bool]: ...
+    def __gt__(self, other: TimestampSeries | Series[Timestamp]) -> Series[bool]: ...
     # error: Signature of "__add__" incompatible with supertype "date"/"datetime"
     @overload  # type: ignore[override]
-    def __add__(self, other: np.ndarray) -> np.ndarray: ...
+    def __add__(
+        self, other: npt.NDArray[np.timedelta64]
+    ) -> npt.NDArray[np.datetime64]: ...
     @overload
     def __add__(
         self: _DatetimeT, other: timedelta | np.timedelta64 | Tick
     ) -> _DatetimeT: ...
+    @overload
+    def __add__(
+        self, other: TimedeltaSeries | Series[Timedelta]
+    ) -> TimestampSeries: ...
+    @overload
+    def __add__(self, other: TimedeltaIndex) -> DatetimeIndex: ...
+    @overload
     def __radd__(self: _DatetimeT, other: timedelta) -> _DatetimeT: ...
+    @overload
+    def __radd__(self, other: TimedeltaIndex) -> DatetimeIndex: ...
+    @overload
+    def __radd__(
+        self, other: npt.NDArray[np.timedelta64]
+    ) -> npt.NDArray[np.datetime64]: ...
+    # TODO: test dt64
     @overload  # type: ignore[override]
-    def __sub__(self, other: datetime) -> Timedelta: ...
+    def __sub__(self, other: Timestamp | datetime | np.datetime64) -> Timedelta: ...
     @overload
     def __sub__(
         self: _DatetimeT, other: timedelta | np.timedelta64 | Tick
     ) -> _DatetimeT: ...
-    def __hash__(self) -> int: ...
+    @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]: ...
+    @overload
+    def __eq__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self, other: TimestampSeries | Series[Timestamp]) -> Series[bool]: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self, other: object) -> Literal[False]: ...
+    @overload
+    def __ne__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self, other: TimestampSeries | Series[Timestamp]) -> Series[bool]: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self, other: object) -> Literal[True]: ...
     def weekday(self) -> int: ...
     def isoweekday(self) -> int: ...
     def isocalendar(self) -> tuple[int, int, int]: ...
@@ -198,19 +261,28 @@ class Timestamp(datetime):
     def tz_localize(
         self: _DatetimeT,
         tz: _tzinfo | str | None,
-        ambiguous: str = ...,
-        nonexistent: str = ...,
+        ambiguous: _Ambiguous = ...,
+        nonexistent: _Nonexistent = ...,
     ) -> _DatetimeT: ...
     def normalize(self: _DatetimeT) -> _DatetimeT: ...
     # TODO: round/floor/ceil could return NaT?
     def round(
-        self: _DatetimeT, freq: str, ambiguous: bool | str = ..., nonexistent: str = ...
+        self: _DatetimeT,
+        freq: str,
+        ambiguous: _Ambiguous = ...,
+        nonexistent: _Nonexistent = ...,
     ) -> _DatetimeT: ...
     def floor(
-        self: _DatetimeT, freq: str, ambiguous: bool | str = ..., nonexistent: str = ...
+        self: _DatetimeT,
+        freq: str,
+        ambiguous: _Ambiguous = ...,
+        nonexistent: _Nonexistent = ...,
     ) -> _DatetimeT: ...
     def ceil(
-        self: _DatetimeT, freq: str, ambiguous: bool | str = ..., nonexistent: str = ...
+        self: _DatetimeT,
+        freq: str,
+        ambiguous: _Ambiguous = ...,
+        nonexistent: _Nonexistent = ...,
     ) -> _DatetimeT: ...
     def day_name(self, locale: str | None = ...) -> str: ...
     def month_name(self, locale: str | None = ...) -> str: ...
@@ -223,12 +295,12 @@ class Timestamp(datetime):
     @property
     def dayofyear(self) -> int: ...
     @property
+    def weekofyear(self) -> int: ...
+    @property
     def quarter(self) -> int: ...
     @property
     def week(self) -> int: ...
-    def to_numpy(
-        self, dtype: np.dtype | None = ..., copy: bool = ...
-    ) -> np.datetime64: ...
+    def to_numpy(self) -> np.datetime64: ...
     @property
     def days_in_month(self) -> int: ...
     @property

pyproject.toml

@@ -58,6 +58,7 @@ tabulate = ">=0.8.10"
 jinja2 = "^3.1"
 scipy = ">=1.9.1"
 SQLAlchemy = "^1.4.41"
+types-python-dateutil = ">=2.8.19"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]