Make datetime generic over tzinfo

stdlib/@tests/test_cases/check_datetime.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from datetime import date, datetime, time, timedelta, timezone, tzinfo
+from typing import Any, Union, cast
+from typing_extensions import Never, assert_type
+
+UTC: timezone = timezone.utc
+
+dt_none = cast(datetime[None], None)
+dt_tz = cast(datetime[tzinfo], None)
+dt_both = cast(datetime[Union[tzinfo, None]], None)
+
+# Constructors
+
+assert_type(datetime(2000, 1, 1), datetime[None])
+assert_type(datetime(2000, 1, 1, tzinfo=None), datetime[None])
+assert_type(datetime(2000, 1, 1, tzinfo=UTC), datetime[tzinfo])
+
+assert_type(datetime.fromtimestamp(0), datetime[None])
+assert_type(datetime.fromtimestamp(0, None), datetime[None])
+assert_type(datetime.fromtimestamp(0, UTC), datetime[tzinfo])
+assert_type(datetime.utcfromtimestamp(0), datetime[None])  # pyright: ignore[reportDeprecated]
+
+assert_type(datetime.now(), datetime[None])
+assert_type(datetime.now(None), datetime[None])
+assert_type(datetime.now(UTC), datetime[tzinfo])
+assert_type(datetime.today(), datetime[None])
+assert_type(datetime.utcnow(), datetime[None])  # pyright: ignore[reportDeprecated]
+
+assert_type(datetime.fromisoformat("2000-01-01"), Union[datetime[tzinfo], datetime[None]])
+
+assert_type(datetime.strptime("", ""), datetime[Any])
+
+# Comparisons
+
+assert_type(dt_none < dt_none, bool)
+assert_type(dt_tz < dt_tz, bool)
+assert_type(dt_both < dt_both, bool)
+
+assert_type(dt_none < dt_tz, Never)
+assert_type(dt_tz < dt_none, Never)
+assert_type(dt_both < dt_none, bool)
+assert_type(dt_both < dt_tz, bool)
+assert_type(dt_none < dt_both, bool)
+
+# Sub
+
+assert_type(dt_none - dt_none, timedelta)
+assert_type(dt_tz - dt_tz, timedelta)
+assert_type(dt_both - dt_both, timedelta)
+
+assert_type(dt_none - dt_tz, Never)
+assert_type(dt_tz - dt_none, Never)
+assert_type(dt_both - dt_none, timedelta)
+assert_type(dt_both - dt_tz, timedelta)
+assert_type(dt_none - dt_both, timedelta)
+assert_type(dt_tz - dt_both, timedelta)
+
+# Combine
+
+assert_type(datetime.combine(date(2000, 1, 1), time(12, 0)), datetime[None])
+assert_type(datetime.combine(date(2000, 1, 1), time(12, 0), tzinfo=None), datetime[None])
+assert_type(datetime.combine(date(2000, 1, 1), time(12, 0), tzinfo=UTC), datetime[tzinfo])
+
+# Replace
+
+assert_type(dt_none.replace(year=2001), datetime[None])
+assert_type(dt_none.replace(year=2001, tzinfo=None), datetime[None])
+assert_type(dt_none.replace(year=2001, tzinfo=UTC), datetime[tzinfo])
+assert_type(dt_tz.replace(year=2001), datetime[tzinfo])
+assert_type(dt_tz.replace(year=2001, tzinfo=None), datetime[None])
+assert_type(dt_tz.replace(year=2001, tzinfo=UTC), datetime[tzinfo])
+assert_type(dt_both.replace(year=2001), datetime[Union[tzinfo, None]])
+assert_type(dt_both.replace(year=2001, tzinfo=None), datetime[None])
+assert_type(dt_both.replace(year=2001, tzinfo=UTC), datetime[tzinfo])
+
+# Attributes
+
+assert_type(dt_none.tzinfo, None)
+assert_type(dt_tz.tzinfo, tzinfo)
+assert_type(dt_both.tzinfo, Union[tzinfo, None])

stdlib/@tests/test_cases/check_unittest.py

@@ -19,6 +19,7 @@
 case.assertAlmostEqual(2.4, 2.41)
 case.assertAlmostEqual(Fraction(49, 50), Fraction(48, 50))
 case.assertAlmostEqual(3.14, complex(5, 6))
+case.assertAlmostEqual(datetime(1999, 1, 2), datetime(1999, 1, 2, microsecond=1))
 case.assertAlmostEqual(datetime(1999, 1, 2), datetime(1999, 1, 2, microsecond=1), delta=timedelta(hours=1))
 case.assertAlmostEqual(datetime(1999, 1, 2), datetime(1999, 1, 2, microsecond=1), None, "foo", timedelta(hours=1))
 case.assertAlmostEqual(Decimal("1.1"), Decimal("1.11"))
@@ -28,7 +29,6 @@
 
 case.assertAlmostEqual(2.4, 2.41, places=9, delta=0.02)  # type: ignore
 case.assertAlmostEqual("foo", "bar")  # type: ignore
-case.assertAlmostEqual(datetime(1999, 1, 2), datetime(1999, 1, 2, microsecond=1))  # type: ignore
 case.assertAlmostEqual(Decimal("0.4"), Fraction(1, 2))  # type: ignore
 case.assertAlmostEqual(complex(2, 3), Decimal("0.9"))  # type: ignore
 
@@ -39,12 +39,12 @@
 case.assertAlmostEqual(1, 2.4)
 case.assertNotAlmostEqual(Fraction(49, 50), Fraction(48, 50))
 case.assertAlmostEqual(3.14, complex(5, 6))
+case.assertNotAlmostEqual(datetime(1999, 1, 2), datetime(1999, 1, 2, microsecond=1))
 case.assertNotAlmostEqual(datetime(1999, 1, 2), datetime(1999, 1, 2, microsecond=1), delta=timedelta(hours=1))
 case.assertNotAlmostEqual(datetime(1999, 1, 2), datetime(1999, 1, 2, microsecond=1), None, "foo", timedelta(hours=1))
 
 case.assertNotAlmostEqual(2.4, 2.41, places=9, delta=0.02)  # type: ignore
 case.assertNotAlmostEqual("foo", "bar")  # type: ignore
-case.assertNotAlmostEqual(datetime(1999, 1, 2), datetime(1999, 1, 2, microsecond=1))  # type: ignore
 case.assertNotAlmostEqual(Decimal("0.4"), Fraction(1, 2))  # type: ignore
 case.assertNotAlmostEqual(complex(2, 3), Decimal("0.9"))  # type: ignore
 

stdlib/datetime.pyi

@@ -1,7 +1,7 @@
 import sys
 from abc import abstractmethod
 from time import struct_time
-from typing import ClassVar, Final, SupportsIndex, TypeAlias, final, overload, type_check_only
+from typing import Any, ClassVar, Final, Generic, SupportsIndex, TypeAlias, TypeVar, final, overload, type_check_only
 from typing_extensions import CapsuleType, Never, Self, deprecated, disjoint_base
 
 if sys.version_info >= (3, 11):
@@ -266,10 +266,14 @@ class timedelta:
     def __bool__(self) -> bool: ...
     def __hash__(self) -> int: ...
 
+_TzInfoT = TypeVar("_TzInfoT", bound=tzinfo | None, default=Any)
+
 @disjoint_base
-class datetime(date):
-    min: ClassVar[datetime]
-    max: ClassVar[datetime]
+class datetime(date, Generic[_TzInfoT]):
+    min: ClassVar[datetime[None]]
+    max: ClassVar[datetime[None]]
+
+    @overload
     def __new__(
         cls,
         year: SupportsIndex,
@@ -279,10 +283,39 @@ class datetime(date):
         minute: SupportsIndex = 0,
         second: SupportsIndex = 0,
         microsecond: SupportsIndex = 0,
-        tzinfo: _TzInfo | None = None,
+        tzinfo: None = None,
         *,
         fold: int = 0,
-    ) -> Self: ...
+    ) -> datetime[None]: ...
+    @overload
+    def __new__(
+        cls,
+        year: SupportsIndex,
+        month: SupportsIndex,
+        day: SupportsIndex,
+        hour: SupportsIndex = 0,
+        minute: SupportsIndex = 0,
+        second: SupportsIndex = 0,
+        microsecond: SupportsIndex = 0,
+        *,
+        tzinfo: _TzInfo,
+        fold: int = 0,
+    ) -> datetime[_TzInfo]: ...
+    @overload
+    def __new__(
+        cls,
+        year: SupportsIndex,
+        month: SupportsIndex,
+        day: SupportsIndex,
+        hour: SupportsIndex,
+        minute: SupportsIndex,
+        second: SupportsIndex,
+        microsecond: SupportsIndex,
+        tzinfo: _TzInfo,
+        *,
+        fold: int = 0,
+    ) -> datetime[_TzInfo]: ...
+
     @property
     def hour(self) -> int: ...
     @property
@@ -292,32 +325,64 @@ class datetime(date):
     @property
     def microsecond(self) -> int: ...
     @property
-    def tzinfo(self) -> _TzInfo | None: ...
+    def tzinfo(self) -> _TzInfoT: ...
     @property
     def fold(self) -> int: ...
     # On <3.12, the name of the first parameter in the pure-Python implementation
     # didn't match the name in the C implementation,
     # meaning it is only *safe* to pass it as a keyword argument on 3.12+
     if sys.version_info >= (3, 12):
+        @overload  # type: ignore[override]
+        @classmethod
+        def fromtimestamp(cls, timestamp: float, tz: None = None) -> datetime[None]: ...
+        @overload
         @classmethod
-        def fromtimestamp(cls, timestamp: float, tz: _TzInfo | None = None) -> Self: ...
+        def fromtimestamp(cls, timestamp: float, tz: _TzInfo) -> datetime[_TzInfo]: ...
+        @overload
+        @classmethod
+        def fromtimestamp(cls, timestamp: float, tz: _TzInfo | None) -> datetime[_TzInfo] | datetime[None]: ...
+    else:
+        @overload  # type: ignore[override]
+        @classmethod
+        def fromtimestamp(cls, timestamp: float, /, tz: None = None) -> datetime[None]: ...
+        @overload
+        @classmethod
+        def fromtimestamp(cls, timestamp: float, /, tz: _TzInfo) -> datetime[_TzInfo]: ...
+        @overload
+        @classmethod
+        def fromtimestamp(cls, timestamp: float, /, tz: _TzInfo | None) -> datetime[_TzInfo] | datetime[None]: ...
+
+    # fromisoformat returns Self, but we can't specialize "Self"
+    if sys.version_info >= (3, 15):
+        @classmethod
+        def fromisoformat(cls, string: str, /) -> datetime[_TzInfo] | datetime[None]: ...  # type: ignore[override]
     else:
         @classmethod
-        def fromtimestamp(cls, timestamp: float, /, tz: _TzInfo | None = None) -> Self: ...
+        def fromisoformat(cls, date_string: str, /) -> datetime[_TzInfo] | datetime[None]: ...  # type: ignore[override]
 
     @classmethod
-    @deprecated("Use timezone-aware objects to represent datetimes in UTC; e.g. by calling .fromtimestamp(datetime.timezone.utc)")
-    def utcfromtimestamp(cls, t: float, /) -> Self: ...
+    @deprecated("Use timezone-aware objects to represent datetimes in UTC; e.g. by calling .fromtimestamp(datetime.UTC)")
+    def utcfromtimestamp(cls, t: float, /) -> datetime[None]: ...
+
+    @overload
     @classmethod
-    def now(cls, tz: _TzInfo | None = None) -> Self: ...
+    def now(cls, tz: None = None) -> datetime[None]: ...
+    @overload
     @classmethod
-    @deprecated("Use timezone-aware objects to represent datetimes in UTC; e.g. by calling .now(datetime.timezone.utc)")
-    def utcnow(cls) -> Self: ...
+    def now(cls, tz: _TzInfo) -> datetime[_TzInfo]: ...
+
     @classmethod
-    def combine(cls, date: _Date, time: _Time, tzinfo: _TzInfo | None = ...) -> Self: ...
-    if sys.version_info >= (3, 15):
-        @classmethod
-        def fromisoformat(cls, string: str, /) -> Self: ...
+    def today(cls) -> datetime[None]: ...  # type: ignore[override]
+    @classmethod
+    @deprecated("Use timezone-aware objects to represent datetimes in UTC; e.g. by calling .now(datetime.UTC)")
+    def utcnow(cls) -> datetime[None]: ...
+
+    @overload
+    @classmethod
+    def combine(cls, date: _Date, time: _Time, tzinfo: None = None) -> datetime[None]: ...
+    @overload
+    @classmethod
+    def combine(cls, date: _Date, time: _Time, tzinfo: _TzInfo) -> datetime[_TzInfo]: ...
 
     def timestamp(self) -> float: ...
     def utctimetuple(self) -> struct_time: ...
@@ -340,6 +405,7 @@ class datetime(date):
             fold: int = ...,
         ) -> Self: ...
 
+    @overload
     def replace(
         self,
         year: SupportsIndex = ...,
@@ -349,11 +415,67 @@ class datetime(date):
         minute: SupportsIndex = ...,
         second: SupportsIndex = ...,
         microsecond: SupportsIndex = ...,
-        tzinfo: _TzInfo | None = ...,
         *,
         fold: int = ...,
     ) -> Self: ...
-    def astimezone(self, tz: _TzInfo | None = None) -> Self: ...
+    @overload
+    def replace(
+        self,
+        year: SupportsIndex,
+        month: SupportsIndex,
+        day: SupportsIndex,
+        hour: SupportsIndex,
+        minute: SupportsIndex,
+        second: SupportsIndex,
+        microsecond: SupportsIndex,
+        tzinfo: _TzInfo,
+        *,
+        fold: int = ...,
+    ) -> datetime[_TzInfo]: ...
+    @overload
+    def replace(
+        self,
+        year: SupportsIndex = ...,
+        month: SupportsIndex = ...,
+        day: SupportsIndex = ...,
+        hour: SupportsIndex = ...,
+        minute: SupportsIndex = ...,
+        second: SupportsIndex = ...,
+        microsecond: SupportsIndex = ...,
+        *,
+        tzinfo: _TzInfo,
+        fold: int = ...,
+    ) -> datetime[_TzInfo]: ...
+    @overload
+    def replace(
+        self,
+        year: SupportsIndex,
+        month: SupportsIndex,
+        day: SupportsIndex,
+        hour: SupportsIndex,
+        minute: SupportsIndex,
+        second: SupportsIndex,
+        microsecond: SupportsIndex,
+        tzinfo: None,
+        *,
+        fold: int = ...,
+    ) -> datetime[None]: ...
+    @overload
+    def replace(
+        self,
+        year: SupportsIndex = ...,
+        month: SupportsIndex = ...,
+        day: SupportsIndex = ...,
+        hour: SupportsIndex = ...,
+        minute: SupportsIndex = ...,
+        second: SupportsIndex = ...,
+        microsecond: SupportsIndex = ...,
+        *,
+        tzinfo: None,
+        fold: int = ...,
+    ) -> datetime[None]: ...
+
+    def astimezone(self, tz: _TzInfo | None = None) -> datetime[_TzInfo]: ...
     def isoformat(self, sep: str = "T", timespec: str = "auto") -> str: ...
 
     if sys.version_info >= (3, 15):
@@ -366,15 +488,56 @@ class datetime(date):
     def utcoffset(self) -> timedelta | None: ...
     def tzname(self) -> str | None: ...
     def dst(self) -> timedelta | None: ...
-    def __le__(self, value: datetime, /) -> bool: ...  # type: ignore[override]
-    def __lt__(self, value: datetime, /) -> bool: ...  # type: ignore[override]
-    def __ge__(self, value: datetime, /) -> bool: ...  # type: ignore[override]
-    def __gt__(self, value: datetime, /) -> bool: ...  # type: ignore[override]
+
+    @overload  # type: ignore[override]
+    def __le__(  # type: ignore[overload-overlap]
+        self: datetime[_TzInfo] | datetime[_TzInfo | None], value: datetime[_TzInfo] | datetime[_TzInfo | None], /
+    ) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __le__(self: datetime[None] | datetime[_TzInfo | None], value: datetime[None] | datetime[_TzInfo | None], /) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __le__(self: datetime[Any], value: datetime[Any], /) -> Never: ...
+
+    @overload  # type: ignore[override]
+    def __lt__(  # type: ignore[overload-overlap]
+        self: datetime[_TzInfo] | datetime[_TzInfo | None], value: datetime[_TzInfo] | datetime[_TzInfo | None], /
+    ) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __lt__(self: datetime[None] | datetime[_TzInfo | None], value: datetime[None] | datetime[_TzInfo | None], /) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __lt__(self: datetime[Any], value: datetime[Any], /) -> Never: ...
+
+    @overload  # type: ignore[override]
+    def __ge__(  # type: ignore[overload-overlap]
+        self: datetime[_TzInfo] | datetime[_TzInfo | None], value: datetime[_TzInfo] | datetime[_TzInfo | None], /
+    ) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __ge__(self: datetime[None] | datetime[_TzInfo | None], value: datetime[None] | datetime[_TzInfo | None], /) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __ge__(self: datetime[Any], value: datetime[Any], /) -> Never: ...
+
+    @overload  # type: ignore[override]
+    def __gt__(  # type: ignore[overload-overlap]
+        self: datetime[_TzInfo] | datetime[_TzInfo | None], value: datetime[_TzInfo] | datetime[_TzInfo | None], /
+    ) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __gt__(self: datetime[None] | datetime[_TzInfo | None], value: datetime[None] | datetime[_TzInfo | None], /) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __gt__(self: datetime[Any], value: datetime[Any], /) -> Never: ...
+
     def __eq__(self, value: object, /) -> bool: ...
     def __hash__(self) -> int: ...
 
     @overload  # type: ignore[override]
-    def __sub__(self, value: Self, /) -> timedelta: ...
+    def __sub__(  # type: ignore[overload-overlap]
+        self: datetime[_TzInfo] | datetime[_TzInfo | None], value: datetime[_TzInfo] | datetime[_TzInfo | None], /
+    ) -> timedelta: ...
+    @overload
+    def __sub__(  # type: ignore[overload-overlap]
+        self: datetime[None] | datetime[_TzInfo | None], value: datetime[None] | datetime[_TzInfo | None], /
+    ) -> timedelta: ...
+    @overload
+    def __sub__(self: datetime[Any], value: datetime[Any], /) -> Never: ...
     @overload
     def __sub__(self, value: timedelta, /) -> Self: ...