Improve Babel dates stub (#6730)

stubs/babel/METADATA.toml

@@ -1 +1,2 @@
-version = "2.9.*"
+version = "2.9.*"
+requires = ["types-pytz"]

stubs/babel/babel/dates.pyi

@@ -1,30 +1,100 @@
-from datetime import date, datetime, time
-from typing import Any
+from datetime import date, datetime, time, timedelta, tzinfo
+from typing import Any, overload
+from typing_extensions import Literal
 
-NO_INHERITANCE_MARKER: str
-LC_TIME: Any
-date_ = date
-datetime_ = datetime
-time_ = time
+from babel.core import Locale
+from babel.util import LOCALTZ as LOCALTZ, UTC as UTC
+from pytz import BaseTzInfo
 
-def get_timezone(zone: Any | None = ...): ...
-def get_next_timezone_transition(zone: Any | None = ..., dt: Any | None = ...): ...
+# The module contents here are organized the same way they are in the API documentation at
+# http://babel.pocoo.org/en/latest/api/dates.html
+
+# Date and Time Formatting
+_Instant = date | time | datetime | float | None
+_PredefinedTimeFormat = Literal["full", "long", "medium", "short"]
+
+def format_datetime(
+    datetime: _Instant = ..., format: _PredefinedTimeFormat | str = ..., tzinfo: tzinfo | None = ..., locale: str | Locale = ...
+) -> str: ...
+def format_date(
+    date: date | datetime | None = ..., format: _PredefinedTimeFormat | str = ..., locale: str | Locale = ...
+) -> str: ...
+def format_time(
+    time: time | datetime | float | None = ...,
+    format: _PredefinedTimeFormat | str = ...,
+    tzinfo: tzinfo | None = ...,
+    locale: str | Locale = ...,
+) -> str: ...
+def format_timedelta(
+    delta: timedelta | int,
+    granularity: Literal["year", "month", "week", "day", "hour", "minute", "second"] = ...,
+    threshold: float = ...,
+    add_direction: bool = ...,
+    format: Literal["narrow", "short", "medium", "long"] = ...,
+    locale: str | Locale = ...,
+) -> str: ...
+def format_skeleton(
+    skeleton: str, datetime: _Instant = ..., tzinfo: tzinfo | None = ..., fuzzy: bool = ..., locale: str | Locale = ...
+) -> str: ...
+def format_interval(
+    start: _Instant,
+    end: _Instant,
+    skeleton: str | None = ...,
+    tzinfo: tzinfo | None = ...,
+    fuzzy: bool = ...,
+    locale: str | Locale = ...,
+) -> str: ...
+
+# Timezone Functionality
+@overload
+def get_timezone(zone: str | BaseTzInfo | None = ...) -> BaseTzInfo: ...
+@overload
+def get_timezone(zone: tzinfo) -> tzinfo: ...
+def get_timezone_gmt(
+    datetime: _Instant = ...,
+    width: Literal["long", "short", "iso8601", "iso8601_short"] = ...,
+    locale: str | Locale = ...,
+    return_z: bool = ...,
+) -> str: ...
+
+_DtOrTzinfo = datetime | tzinfo | str | int | time | None
+
+def get_timezone_location(dt_or_tzinfo: _DtOrTzinfo = ..., locale: str | Locale = ..., return_city: bool = ...) -> str: ...
+def get_timezone_name(
+    dt_or_tzinfo: _DtOrTzinfo = ...,
+    width: Literal["long", "short"] = ...,
+    uncommon: bool = ...,
+    locale: str | Locale = ...,
+    zone_variant: Literal["generic", "daylight", "standard"] | None = ...,
+    return_zone: bool = ...,
+) -> str: ...
+
+# Note: While Babel accepts any tzinfo for the most part, the get_next_timeout_transition()
+# function requires a tzinfo that is produced by get_timezone()/pytz AND has DST info.
+# The typing here will help you with the first requirement, but will not protect against
+# pytz tzinfo's without DST info, like what you get from get_timezone("UTC") for instance.
+def get_next_timezone_transition(zone: BaseTzInfo | None = ..., dt: _Instant = ...) -> TimezoneTransition: ...
 
 class TimezoneTransition:
-    activates: Any
-    from_tzinfo: Any
-    to_tzinfo: Any
-    reference_date: Any
-    def __init__(self, activates, from_tzinfo, to_tzinfo, reference_date: Any | None = ...) -> None: ...
+    # This class itself is not included in the documentation, yet it is mentioned by name.
+    # See https://github.com/python-babel/babel/issues/823
+    activates: datetime
+    from_tzinfo: tzinfo
+    to_tzinfo: tzinfo
+    reference_date: datetime | None
+    def __init__(
+        self, activates: datetime, from_tzinfo: tzinfo, to_tzinfo: tzinfo, reference_date: datetime | None = ...
+    ) -> None: ...
     @property
-    def from_tz(self): ...
+    def from_tz(self) -> str: ...
     @property
-    def to_tz(self): ...
+    def to_tz(self) -> str: ...
     @property
-    def from_offset(self): ...
+    def from_offset(self) -> int: ...
     @property
-    def to_offset(self): ...
+    def to_offset(self) -> int: ...
 
+# Data Access
 def get_period_names(width: str = ..., context: str = ..., locale=...): ...
 def get_day_names(width: str = ..., context: str = ..., locale=...): ...
 def get_month_names(width: str = ..., context: str = ..., locale=...): ...
@@ -33,30 +103,22 @@ def get_era_names(width: str = ..., locale=...): ...
 def get_date_format(format: str = ..., locale=...): ...
 def get_datetime_format(format: str = ..., locale=...): ...
 def get_time_format(format: str = ..., locale=...): ...
-def get_timezone_gmt(datetime: Any | None = ..., width: str = ..., locale=..., return_z: bool = ...): ...
-def get_timezone_location(dt_or_tzinfo: Any | None = ..., locale=..., return_city: bool = ...): ...
-def get_timezone_name(
-    dt_or_tzinfo: Any | None = ...,
-    width: str = ...,
-    uncommon: bool = ...,
-    locale=...,
-    zone_variant: Any | None = ...,
-    return_zone: bool = ...,
-): ...
-def format_date(date: Any | None = ..., format: str = ..., locale=...): ...
-def format_datetime(datetime: Any | None = ..., format: str = ..., tzinfo: Any | None = ..., locale=...): ...
-def format_time(time: Any | None = ..., format: str = ..., tzinfo: Any | None = ..., locale=...): ...
-def format_skeleton(skeleton, datetime: Any | None = ..., tzinfo: Any | None = ..., fuzzy: bool = ..., locale=...): ...
+
+# Basic Parsing
+def parse_date(string, locale=...): ...
+def parse_time(string, locale=...): ...
+def parse_pattern(pattern): ...
+
+# Undocumented
+NO_INHERITANCE_MARKER: str
+LC_TIME: Any
+date_ = date
+datetime_ = datetime
+time_ = time
 
 TIMEDELTA_UNITS: Any
 
-def format_timedelta(
-    delta, granularity: str = ..., threshold: float = ..., add_direction: bool = ..., format: str = ..., locale=...
-): ...
-def format_interval(start, end, skeleton: Any | None = ..., tzinfo: Any | None = ..., fuzzy: bool = ..., locale=...): ...
 def get_period_id(time, tzinfo: Any | None = ..., type: Any | None = ..., locale=...): ...
-def parse_date(string, locale=...): ...
-def parse_time(string, locale=...): ...
 
 class DateTimePattern:
     pattern: Any
@@ -91,7 +153,6 @@ class DateTimeFormat:
 PATTERN_CHARS: Any
 PATTERN_CHAR_ORDER: str
 
-def parse_pattern(pattern): ...
 def tokenize_pattern(pattern): ...
 def untokenize_pattern(tokens): ...
 def split_interval_pattern(pattern): ...

stubs/babel/babel/util.pyi

@@ -4,6 +4,7 @@ from datetime import tzinfo
 from typing import Any
 
 from babel import localtime as localtime
+from pytz import BaseTzInfo
 
 missing: Any
 
@@ -32,8 +33,8 @@ class FixedOffsetTimezone(tzinfo):
     def tzname(self, dt): ...
     def dst(self, dt): ...
 
-UTC: Any
-LOCALTZ: Any
+UTC: BaseTzInfo
+LOCALTZ: BaseTzInfo
 get_localzone = localtime.get_localzone
 STDOFFSET: Any
 DSTOFFSET: Any