From e39fa2f93c7b16fb03a753bfd91e238d083fe403 Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Tue, 8 Oct 2024 10:27:20 +0200 Subject: [PATCH 1/3] more precise NaT stub --- pandas/_libs/tslibs/nattype.pyi | 39 +++++++++++++++------------------ 1 file changed, 18 insertions(+), 21 deletions(-) diff --git a/pandas/_libs/tslibs/nattype.pyi b/pandas/_libs/tslibs/nattype.pyi index fa1577f033fff..75d81cc05b89d 100644 --- a/pandas/_libs/tslibs/nattype.pyi +++ b/pandas/_libs/tslibs/nattype.pyi @@ -25,12 +25,9 @@ NaT: NaTType iNaT: int nat_strings: set[str] -_NaTComparisonTypes: TypeAlias = ( - datetime | timedelta | Period | np.datetime64 | np.timedelta64 -) +_TimeLike: TypeAlias = datetime | timedelta | Period | np.datetime64 | np.timedelta64 +_TimeDelta: TypeAlias = timedelta | np.timedelta64 -class _NatComparison: - def __call__(self, other: _NaTComparisonTypes) -> bool: ... class NaTType: _value: np.int64 @@ -161,30 +158,30 @@ class NaTType: @property def qyear(self) -> float: ... # comparisons - def __eq__(self, other: object) -> bool: ... - def __ne__(self, other: object) -> bool: ... - __lt__: _NatComparison - __le__: _NatComparison - __gt__: _NatComparison - __ge__: _NatComparison + def __eq__(self, other: object, /) -> bool: ... + def __ne__(self, other: object, /) -> bool: ... + def __lt__(self, other: Self | _TimeLike, /) -> Literal[False]: ... + def __le__(self, other: Self | _TimeLike, /) -> Literal[False]: ... + def __gt__(self, other: Self | _TimeLike, /) -> Literal[False]: ... + def __ge__(self, other: Self | _TimeLike, /) -> Literal[False]: ... # unary operators def __pos__(self) -> Self: ... def __neg__(self) -> Self: ... # binary operators - def __sub__(self, other: Self | timedelta | datetime) -> Self: ... - def __rsub__(self, other: Self | timedelta | datetime) -> Self: ... - def __add__(self, other: Self | timedelta | datetime) -> Self: ... - def __radd__(self, other: Self | timedelta | datetime) -> Self: ... - def __mul__(self, other: float) -> Self: ... # analogous to timedelta - def __rmul__(self, other: float) -> Self: ... + def __sub__(self, other: Self | _TimeLike, /) -> Self: ... + def __rsub__(self, other: Self | _TimeLike, /) -> Self: ... + def __add__(self, other: Self | _TimeLike, /) -> Self: ... + def __radd__(self, other: Self | _TimeLike, /) -> Self: ... + def __mul__(self, other: float, /) -> Self: ... # analogous to timedelta + def __rmul__(self, other: float, /) -> Self: ... @overload # analogous to timedelta - def __truediv__(self, other: Self | timedelta) -> float: ... # Literal[NaN] + def __truediv__(self, other: Self | _TimeDelta, /) -> float: ... # Literal[NaN] @overload - def __truediv__(self, other: float) -> Self: ... + def __truediv__(self, other: float, /) -> Self: ... @overload # analogous to timedelta - def __floordiv__(self, other: Self | timedelta) -> float: ... # Literal[NaN] + def __floordiv__(self, other: Self | _TimeDelta, /) -> float: ... # Literal[NaN] @overload - def __floordiv__(self, other: float) -> Self: ... + def __floordiv__(self, other: float, /) -> Self: ... # other def __hash__(self) -> int: ... def as_unit(self, unit: str, round_ok: bool = ...) -> NaTType: ... From 115f31d3a9a034def1fa94a7090a87e107fa16e6 Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Tue, 8 Oct 2024 10:36:51 +0200 Subject: [PATCH 2/3] ruff format --- pandas/_libs/tslibs/nattype.pyi | 1 - 1 file changed, 1 deletion(-) diff --git a/pandas/_libs/tslibs/nattype.pyi b/pandas/_libs/tslibs/nattype.pyi index 75d81cc05b89d..536ce4e94301f 100644 --- a/pandas/_libs/tslibs/nattype.pyi +++ b/pandas/_libs/tslibs/nattype.pyi @@ -28,7 +28,6 @@ nat_strings: set[str] _TimeLike: TypeAlias = datetime | timedelta | Period | np.datetime64 | np.timedelta64 _TimeDelta: TypeAlias = timedelta | np.timedelta64 - class NaTType: _value: np.int64 @property From 4a595c7146ece6b4825a208ea54059c9f9187782 Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Tue, 8 Oct 2024 10:40:54 +0200 Subject: [PATCH 3/3] updated == and != to return literal --- pandas/_libs/tslibs/nattype.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pandas/_libs/tslibs/nattype.pyi b/pandas/_libs/tslibs/nattype.pyi index 536ce4e94301f..d3b10fbe79cb9 100644 --- a/pandas/_libs/tslibs/nattype.pyi +++ b/pandas/_libs/tslibs/nattype.pyi @@ -157,8 +157,8 @@ class NaTType: @property def qyear(self) -> float: ... # comparisons - def __eq__(self, other: object, /) -> bool: ... - def __ne__(self, other: object, /) -> bool: ... + def __eq__(self, other: object, /) -> Literal[False]: ... + def __ne__(self, other: object, /) -> Literal[True]: ... def __lt__(self, other: Self | _TimeLike, /) -> Literal[False]: ... def __le__(self, other: Self | _TimeLike, /) -> Literal[False]: ... def __gt__(self, other: Self | _TimeLike, /) -> Literal[False]: ...