TYP: remove NaTType as possible result of Timestamp and Timedelta constructor (#46171)

Author: Dr-Irv

pandas/_libs/tslibs/nattype.pyi

@@ -3,7 +3,10 @@ from datetime import (
     timedelta,
     tzinfo as _tzinfo,
 )
-from typing import Any
+from typing import (
+    Any,
+    Union,
+)
 
 import numpy as np
 
@@ -15,7 +18,12 @@ nat_strings: set[str]
 
 def is_null_datetimelike(val: object, inat_is_null: bool = ...) -> bool: ...
 
-class NaTType(datetime):
+_NaTComparisonTypes = Union[datetime, timedelta, Period, np.datetime64, np.timedelta64]
+
+class _NatComparison:
+    def __call__(self, other: _NaTComparisonTypes) -> bool: ...
+
+class NaTType:
     value: np.int64
     def asm8(self) -> np.datetime64: ...
     def to_datetime64(self) -> np.datetime64: ...
@@ -54,26 +62,19 @@ class NaTType(datetime):
     def weekofyear(self) -> float: ...
     def day_name(self) -> float: ...
     def month_name(self) -> float: ...
-    # error: Return type "float" of "weekday" incompatible with return
-    # type "int" in supertype "date"
-    def weekday(self) -> float: ...  # type: ignore[override]
-    # error: Return type "float" of "isoweekday" incompatible with return
-    # type "int" in supertype "date"
-    def isoweekday(self) -> float: ...  # type: ignore[override]
+    def weekday(self) -> float: ...
+    def isoweekday(self) -> float: ...
     def total_seconds(self) -> float: ...
-    # error: Signature of "today" incompatible with supertype "datetime"
-    def today(self, *args, **kwargs) -> NaTType: ...  # type: ignore[override]
-    # error: Signature of "today" incompatible with supertype "datetime"
-    def now(self, *args, **kwargs) -> NaTType: ...  # type: ignore[override]
+    def today(self, *args, **kwargs) -> NaTType: ...
+    def now(self, *args, **kwargs) -> NaTType: ...
     def to_pydatetime(self) -> NaTType: ...
     def date(self) -> NaTType: ...
     def round(self) -> NaTType: ...
     def floor(self) -> NaTType: ...
     def ceil(self) -> NaTType: ...
     def tz_convert(self) -> NaTType: ...
     def tz_localize(self) -> NaTType: ...
-    # error: Signature of "replace" incompatible with supertype "datetime"
-    def replace(  # type: ignore[override]
+    def replace(
         self,
         year: int | None = ...,
         month: int | None = ...,
@@ -86,38 +87,24 @@ class NaTType(datetime):
         tzinfo: _tzinfo | None = ...,
         fold: int | None = ...,
     ) -> NaTType: ...
-    # error: Return type "float" of "year" incompatible with return
-    # type "int" in supertype "date"
     @property
-    def year(self) -> float: ...  # type: ignore[override]
+    def year(self) -> float: ...
     @property
     def quarter(self) -> float: ...
-    # error: Return type "float" of "month" incompatible with return
-    # type "int" in supertype "date"
     @property
-    def month(self) -> float: ...  # type: ignore[override]
-    # error: Return type "float" of "day" incompatible with return
-    # type "int" in supertype "date"
+    def month(self) -> float: ...
     @property
-    def day(self) -> float: ...  # type: ignore[override]
-    # error: Return type "float" of "hour" incompatible with return
-    # type "int" in supertype "date"
+    def day(self) -> float: ...
     @property
-    def hour(self) -> float: ...  # type: ignore[override]
-    # error: Return type "float" of "minute" incompatible with return
-    # type "int" in supertype "date"
+    def hour(self) -> float: ...
     @property
-    def minute(self) -> float: ...  # type: ignore[override]
-    # error: Return type "float" of "second" incompatible with return
-    # type "int" in supertype "date"
+    def minute(self) -> float: ...
     @property
-    def second(self) -> float: ...  # type: ignore[override]
+    def second(self) -> float: ...
     @property
     def millisecond(self) -> float: ...
-    # error: Return type "float" of "microsecond" incompatible with return
-    # type "int" in supertype "date"
     @property
-    def microsecond(self) -> float: ...  # type: ignore[override]
+    def microsecond(self) -> float: ...
     @property
     def nanosecond(self) -> float: ...
     # inject Timedelta properties
@@ -132,24 +119,7 @@ class NaTType(datetime):
     def qyear(self) -> float: ...
     def __eq__(self, other: Any) -> bool: ...
     def __ne__(self, other: Any) -> bool: ...
-    # https://github.com/python/mypy/issues/9015
-    # error: Argument 1 of "__lt__" is incompatible with supertype "date";
-    # supertype defines the argument type as "date"
-    def __lt__(  # type: ignore[override]
-        self, other: datetime | timedelta | Period | np.datetime64 | np.timedelta64
-    ) -> bool: ...
-    # error: Argument 1 of "__le__" is incompatible with supertype "date";
-    # supertype defines the argument type as "date"
-    def __le__(  # type: ignore[override]
-        self, other: datetime | timedelta | Period | np.datetime64 | np.timedelta64
-    ) -> bool: ...
-    # error: Argument 1 of "__gt__" is incompatible with supertype "date";
-    # supertype defines the argument type as "date"
-    def __gt__(  # type: ignore[override]
-        self, other: datetime | timedelta | Period | np.datetime64 | np.timedelta64
-    ) -> bool: ...
-    # error: Argument 1 of "__ge__" is incompatible with supertype "date";
-    # supertype defines the argument type as "date"
-    def __ge__(  # type: ignore[override]
-        self, other: datetime | timedelta | Period | np.datetime64 | np.timedelta64
-    ) -> bool: ...
+    __lt__: _NatComparison
+    __le__: _NatComparison
+    __gt__: _NatComparison
+    __ge__: _NatComparison

pandas/_libs/tslibs/timedeltas.pyi

@@ -7,12 +7,12 @@ from typing import (
 )
 
 import numpy as np
+import numpy.typing as npt
 
 from pandas._libs.tslibs import (
     NaTType,
     Tick,
 )
-from pandas._typing import npt
 
 _S = TypeVar("_S", bound=timedelta)
 
@@ -26,21 +26,22 @@ def array_to_timedelta64(
     errors: str = ...,
 ) -> np.ndarray: ...  # np.ndarray[m8ns]
 def parse_timedelta_unit(unit: str | None) -> str: ...
-def delta_to_nanoseconds(delta: Tick | np.timedelta64 | timedelta | int) -> int: ...
+def delta_to_nanoseconds(delta: np.timedelta64 | timedelta | Tick) -> int: ...
 
 class Timedelta(timedelta):
     min: ClassVar[Timedelta]
     max: ClassVar[Timedelta]
     resolution: ClassVar[Timedelta]
     value: int  # np.int64
-
-    # error: "__new__" must return a class instance (got "Union[Timedelta, NaTType]")
-    def __new__(  # type: ignore[misc]
+    def __new__(
         cls: Type[_S],
         value=...,
         unit: str = ...,
         **kwargs: int | float | np.integer | np.floating,
-    ) -> _S | NaTType: ...
+    ) -> _S: ...
+    # GH 46171
+    # While Timedelta can return pd.NaT, having the constructor return
+    # a Union with NaTType makes things awkward for users of pandas
     @property
     def days(self) -> int: ...
     @property

pandas/_libs/tslibs/timestamps.pyi

@@ -32,9 +32,7 @@ class Timestamp(datetime):
 
     resolution: ClassVar[Timedelta]
     value: int  # np.int64
-
-    # error: "__new__" must return a class instance (got "Union[Timestamp, NaTType]")
-    def __new__(  # type: ignore[misc]
+    def __new__(
         cls: type[_DatetimeT],
         ts_input: int
         | np.integer
@@ -57,7 +55,10 @@ class Timestamp(datetime):
         tzinfo: _tzinfo | None = ...,
         *,
         fold: int | None = ...,
-    ) -> _DatetimeT | NaTType: ...
+    ) -> _DatetimeT: ...
+    # GH 46171
+    # While Timestamp can return pd.NaT, having the constructor return
+    # a Union with NaTType makes things awkward for users of pandas
     def _set_freq(self, freq: BaseOffset | None) -> None: ...
     @property
     def year(self) -> int: ...
@@ -145,9 +146,11 @@ class Timestamp(datetime):
     ) -> _DatetimeT: ...
     def __radd__(self: _DatetimeT, other: timedelta) -> _DatetimeT: ...
     @overload  # type: ignore
-    def __sub__(self, other: datetime) -> timedelta: ...
+    def __sub__(self, other: datetime) -> Timedelta: ...
     @overload
-    def __sub__(self, other: timedelta | np.timedelta64 | Tick) -> datetime: ...
+    def __sub__(
+        self: _DatetimeT, other: timedelta | np.timedelta64 | Tick
+    ) -> _DatetimeT: ...
     def __hash__(self) -> int: ...
     def weekday(self) -> int: ...
     def isoweekday(self) -> int: ...
@@ -206,3 +209,5 @@ class Timestamp(datetime):
     def to_numpy(
         self, dtype: np.dtype | None = ..., copy: bool = ...
     ) -> np.datetime64: ...
+    @property
+    def _date_repr(self) -> str: ...

pandas/_typing.py

@@ -35,6 +35,7 @@
     import numpy.typing as npt
 
     from pandas._libs import (
+        NaTType,
         Period,
         Timedelta,
         Timestamp,
@@ -308,3 +309,7 @@ def closed(self) -> bool:
 # Interval closed type
 
 IntervalClosedType = Literal["left", "right", "both", "neither"]
+
+# datetime and NaTType
+
+DatetimeNaTType = Union[datetime, "NaTType"]

pandas/core/arrays/datetimes.py

@@ -777,7 +777,9 @@ def _add_offset(self, offset) -> DatetimeArray:
     def _sub_datetimelike_scalar(self, other):
         # subtract a datetime from myself, yielding a ndarray[timedelta64[ns]]
         assert isinstance(other, (datetime, np.datetime64))
-        assert other is not NaT
+        # error: Non-overlapping identity check (left operand type: "Union[datetime,
+        # datetime64]", right operand type: "NaTType")  [comparison-overlap]
+        assert other is not NaT  # type: ignore[comparison-overlap]
         other = Timestamp(other)
         # error: Non-overlapping identity check (left operand type: "Timestamp",
         # right operand type: "NaTType")

pandas/io/formats/format.py

@@ -1767,16 +1767,13 @@ def _format_datetime64_dateonly(
     nat_rep: str = "NaT",
     date_format: str | None = None,
 ) -> str:
-    if x is NaT:
+    if isinstance(x, NaTType):
         return nat_rep
 
     if date_format:
         return x.strftime(date_format)
     else:
-        # error: Item "NaTType" of "Union[NaTType, Any]" has no attribute "_date_repr"
-        #  The underlying problem here is that mypy doesn't understand that NaT
-        #  is a singleton, so that the check above excludes it here.
-        return x._date_repr  # type: ignore[union-attr]
+        return x._date_repr
 
 
 def get_format_datetime64(

pandas/io/sas/sas_xport.py

@@ -17,6 +17,7 @@
 import numpy as np
 
 from pandas._typing import (
+    DatetimeNaTType,
     FilePath,
     ReadBuffer,
 )
@@ -139,7 +140,7 @@
 """
 
 
-def _parse_date(datestr: str) -> datetime:
+def _parse_date(datestr: str) -> DatetimeNaTType:
     """Given a date in xport format, return Python date."""
     try:
         # e.g. "16FEB11:10:07:55"

pandas/tests/resample/test_datetime_index.py

@@ -1,12 +1,14 @@
 from datetime import datetime
 from functools import partial
 from io import StringIO
+from typing import List
 
 import numpy as np
 import pytest
 import pytz
 
 from pandas._libs import lib
+from pandas._typing import DatetimeNaTType
 from pandas.errors import UnsupportedFunctionCall
 
 import pandas as pd
@@ -1286,7 +1288,7 @@ def test_resample_consistency():
     tm.assert_series_equal(s10_2, rl)
 
 
-dates1 = [
+dates1: List[DatetimeNaTType] = [
     datetime(2014, 10, 1),
     datetime(2014, 9, 3),
     datetime(2014, 11, 5),
@@ -1295,7 +1297,9 @@ def test_resample_consistency():
     datetime(2014, 7, 15),
 ]
 
-dates2 = dates1[:2] + [pd.NaT] + dates1[2:4] + [pd.NaT] + dates1[4:]
+dates2: List[DatetimeNaTType] = (
+    dates1[:2] + [pd.NaT] + dates1[2:4] + [pd.NaT] + dates1[4:]
+)
 dates3 = [pd.NaT] + dates1 + [pd.NaT]  # type: ignore[operator]