ENH: Improve typing for Timestamp

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

pandas-stubs/core/algorithms.pyi

@@ -4,7 +4,6 @@ from typing import (
 )
 
 import numpy as np
-import pandas as pd
 from pandas import (
     Categorical,
     CategoricalIndex,

pandas-stubs/core/arrays/arrow/dtype.pyi

@@ -1,4 +1,3 @@
-import numpy as np
 import pyarrow as pa
 
 from pandas.core.dtypes.base import StorageExtensionDtype

pandas-stubs/core/arrays/numpy_.pyi

@@ -1,4 +1,3 @@
-import numpy as np
 from numpy.lib.mixins import NDArrayOperatorsMixin
 from pandas.core.arrays.base import (
     ExtensionArray,

pandas-stubs/core/dtypes/common.pyi

@@ -1,6 +1,5 @@
 from typing import Union
 
-import numpy as np
 import pandas as pd
 from typing_extensions import TypeAlias
 

pandas-stubs/core/frame.pyi

@@ -110,7 +110,6 @@ from pandas._typing import (
     num,
 )
 
-from pandas.io.formats import format as fmt
 from pandas.io.formats.style import Styler
 from pandas.plotting import PlotAccessor
 

pandas-stubs/core/groupby/generic.pyi

@@ -15,7 +15,6 @@ from matplotlib.axes import (
     Axes as PlotAxes,
     SubplotBase as AxesSubplot,
 )
-import numpy as np
 from pandas.core.frame import DataFrame
 from pandas.core.generic import NDFrame
 from pandas.core.groupby.groupby import (  # , get_groupby as get_groupby

pandas-stubs/core/indexes/interval.pyi

@@ -1,4 +1,3 @@
-import datetime as dt
 from typing import (
     Any,
     Hashable,

pandas-stubs/core/indexes/numeric.pyi

@@ -1,6 +1,5 @@
 from typing import Iterable
 
-import numpy as np
 from pandas.core.indexes.base import Index
 
 from pandas._typing import np_ndarray_int64

pandas-stubs/core/ops/dispatch.pyi

@@ -1,5 +1,3 @@
-import numpy as np
-
 from pandas.core.dtypes.generic import ABCSeries
 
 def should_extension_dispatch(left: ABCSeries, right) -> bool: ...

pandas-stubs/core/tools/datetimes.pyi

@@ -10,7 +10,6 @@ from typing import (
 )
 
 import numpy as np
-import pandas as pd
 from pandas import (
     Index,
     Timestamp,

pandas-stubs/core/tools/timedeltas.pyi

@@ -4,7 +4,6 @@ from typing import (
     overload,
 )
 
-import pandas as pd
 from pandas import Index
 from pandas.core.indexes.timedeltas import TimedeltaIndex
 from pandas.core.series import (

pandas-stubs/core/window/expanding.pyi

@@ -4,7 +4,6 @@ from typing import (
     overload,
 )
 
-import numpy as np
 from pandas import (
     DataFrame,
     Series,

pandas-stubs/core/window/rolling.pyi

@@ -5,7 +5,6 @@ from typing import (
     overload,
 )
 
-import numpy as np
 from pandas import (
     DataFrame,
     Series,

pandas-stubs/io/clipboards.pyi

@@ -8,7 +8,6 @@ from typing import (
     overload,
 )
 
-import numpy as np
 from pandas.core.frame import DataFrame
 from pandas.core.indexes.base import Index
 from pandas.core.series import Series

pandas-stubs/io/formats/format.pyi

@@ -1,5 +1,3 @@
-import numpy as np
-
 class EngFormatter:
     ENG_PREFIXES = ...
     accuracy = ...

pandas-stubs/io/parsers/readers.pyi

@@ -16,7 +16,6 @@ from pandas.core.frame import DataFrame
 from pandas.core.indexes.base import Index
 from pandas.core.series import Series
 
-import pandas._libs.lib as lib
 from pandas._typing import (
     CompressionOptions,
     CSVEngine,

pandas-stubs/io/pytables.pyi

@@ -7,7 +7,6 @@ from typing import (
     overload,
 )
 
-import numpy as np
 from pandas import (
     DataFrame,
     Series,