type tz in tz_convert, tz_localize, date_range, and bdate_range (#1118)

Author: MarcoGorelli

pandas-stubs/_libs/tslibs/timestamps.pyi

@@ -15,6 +15,7 @@ from typing import (
     overload,
 )
 
+from _typing import TimeZones
 import numpy as np
 from pandas import (
     DatetimeIndex,
@@ -70,7 +71,7 @@ class Timestamp(datetime, SupportsIndex):
         tzinfo: _tzinfo | None = ...,
         *,
         nanosecond: int | None = ...,
-        tz: str | _tzinfo | int | None = ...,
+        tz: TimeZones = ...,
         unit: str | int | None = ...,
         fold: Literal[0, 1] | None = ...,
     ) -> Self: ...
@@ -258,11 +259,11 @@ class Timestamp(datetime, SupportsIndex):
     def to_julian_date(self) -> np.float64: ...
     @property
     def asm8(self) -> np.datetime64: ...
-    def tz_convert(self, tz: _tzinfo | str | None) -> Self: ...
+    def tz_convert(self, tz: TimeZones) -> Self: ...
     # TODO: could return NaT?
     def tz_localize(
         self,
-        tz: _tzinfo | str | None,
+        tz: TimeZones,
         ambiguous: _Ambiguous = ...,
         nonexistent: _Nonexistent = ...,
     ) -> Self: ...

pandas-stubs/_typing.pyi

@@ -8,6 +8,7 @@ from collections.abc import (
     Sequence,
 )
 import datetime
+from datetime import tzinfo
 from os import PathLike
 import sys
 from typing import (
@@ -820,4 +821,6 @@ TimeGrouperOrigin: TypeAlias = (
 ExcelReadEngine: TypeAlias = Literal["xlrd", "openpyxl", "odf", "pyxlsb", "calamine"]
 ExcelWriteEngine: TypeAlias = Literal["openpyxl", "odf", "xlsxwriter"]
 
+TimeZones: TypeAlias = str | tzinfo | None | int
+
 __all__ = ["npt", "type_t"]

pandas-stubs/core/arrays/datetimes.pyi

@@ -1,5 +1,6 @@
 from datetime import tzinfo
 
+from _typing import TimeZones
 import numpy as np
 from pandas.core.arrays.datetimelike import (
     DatelikeOps,
@@ -28,8 +29,10 @@ class DatetimeArray(DatetimeLikeArrayMixin, TimelikeOps, DatelikeOps):
     def __array__(self, dtype=...) -> np.ndarray: ...
     def __iter__(self): ...
     def astype(self, dtype, copy: bool = ...): ...
-    def tz_convert(self, tz): ...
-    def tz_localize(self, tz, ambiguous: str = ..., nonexistent: str = ...): ...
+    def tz_convert(self, tz: TimeZones): ...
+    def tz_localize(
+        self, tz: TimeZones, ambiguous: str = ..., nonexistent: str = ...
+    ): ...
     def to_pydatetime(self): ...
     def normalize(self): ...
     def to_period(self, freq=...): ...

pandas-stubs/core/dtypes/dtypes.pyi

@@ -5,6 +5,7 @@ from typing import (
     TypeVar,
 )
 
+from _typing import TimeZones
 import numpy as np
 from pandas.core.indexes.base import Index
 from pandas.core.series import Series
@@ -41,9 +42,7 @@ class CategoricalDtype(PandasExtensionDtype, ExtensionDtype):
     def ordered(self) -> Ordered: ...
 
 class DatetimeTZDtype(PandasExtensionDtype):
-    def __init__(
-        self, unit: Literal["ns"] = ..., tz: str | int | dt.tzinfo | None = ...
-    ) -> None: ...
+    def __init__(self, unit: Literal["ns"] = ..., tz: TimeZones = ...) -> None: ...
     @property
     def unit(self) -> Literal["ns"]: ...
     @property

pandas-stubs/core/frame.pyi

@@ -19,6 +19,7 @@ from typing import (
     overload,
 )
 
+from _typing import TimeZones
 from matplotlib.axes import Axes as PlotAxes
 import numpy as np
 from pandas import (
@@ -2446,14 +2447,14 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     ) -> Self: ...
     def tz_convert(
         self,
-        tz,
+        tz: TimeZones,
         axis: Axis = ...,
         level: Level | None = ...,
         copy: _bool = ...,
     ) -> Self: ...
     def tz_localize(
         self,
-        tz,
+        tz: TimeZones,
         axis: Axis = ...,
         level: Level | None = ...,
         copy: _bool = ...,

pandas-stubs/core/generic.pyi

@@ -546,16 +546,6 @@ class NDFrame(indexing.IndexingMixin):
     ): ...
     def shift(self, periods=..., freq=..., axis=..., fill_value=...) -> Self: ...
     def truncate(self, before=..., after=..., axis=..., copy: _bool = ...) -> Self: ...
-    def tz_convert(self, tz, axis=..., level=..., copy: _bool = ...) -> Self: ...
-    def tz_localize(
-        self,
-        tz,
-        axis=...,
-        level=...,
-        copy: _bool = ...,
-        ambiguous=...,
-        nonexistent: str = ...,
-    ) -> Self: ...
     def abs(self) -> Self: ...
     def describe(self, percentiles=..., include=..., exclude=...) -> NDFrame: ...
     def pct_change(

pandas-stubs/core/indexes/accessors.pyi

@@ -9,6 +9,7 @@ from typing import (
     TypeVar,
 )
 
+from _typing import TimeZones
 import numpy as np
 import numpy.typing as npt
 from pandas import (
@@ -216,15 +217,15 @@ class _DatetimeLikeNoTZMethods(
     ) -> _DTToPeriodReturnType: ...
     def tz_localize(
         self,
-        tz: tzinfo | str | None,
+        tz: TimeZones,
         ambiguous: Literal["raise", "infer", "NaT"] | np_ndarray_bool = ...,
         nonexistent: (
             Literal["shift_forward", "shift_backward", "NaT", "raise"]
             | timedelta
             | Timedelta
         ) = ...,
     ) -> _DTNormalizeReturnType: ...
-    def tz_convert(self, tz: tzinfo | str | None) -> _DTNormalizeReturnType: ...
+    def tz_convert(self, tz: TimeZones) -> _DTNormalizeReturnType: ...
     def normalize(self) -> _DTNormalizeReturnType: ...
     def strftime(self, date_format: str) -> _DTStrKindReturnType: ...
     def month_name(self, locale: str | None = ...) -> _DTStrKindReturnType: ...

pandas-stubs/core/indexes/datetimes.pyi

@@ -9,6 +9,7 @@ from datetime import (
 )
 from typing import overload
 
+from _typing import TimeZones
 import numpy as np
 from pandas import (
     DataFrame,
@@ -92,7 +93,7 @@ def date_range(
     end: str | DateAndDatetimeLike | None = ...,
     periods: int | None = ...,
     freq: str | timedelta | Timedelta | BaseOffset = ...,
-    tz: str | tzinfo = ...,
+    tz: TimeZones = ...,
     normalize: bool = ...,
     name: Hashable | None = ...,
     inclusive: IntervalClosedType = ...,
@@ -104,7 +105,7 @@ def bdate_range(
     end: str | DateAndDatetimeLike | None = ...,
     periods: int | None = ...,
     freq: str | timedelta | Timedelta | BaseOffset = ...,
-    tz: str | tzinfo = ...,
+    tz: TimeZones = ...,
     normalize: bool = ...,
     name: Hashable | None = ...,
     weekmask: str | None = ...,
@@ -118,7 +119,7 @@ def bdate_range(
     periods: int | None = ...,
     *,
     freq: str | timedelta | Timedelta | BaseOffset,
-    tz: str | tzinfo = ...,
+    tz: TimeZones = ...,
     normalize: bool = ...,
     name: Hashable | None = ...,
     weekmask: str | None = ...,

pandas-stubs/core/series.pyi

@@ -24,6 +24,7 @@ from typing import (
     overload,
 )
 
+from _typing import TimeZones
 from matplotlib.axes import (
     Axes as PlotAxes,
     SubplotBase,
@@ -1550,14 +1551,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series[S1]: ...
     def tz_convert(
         self,
-        tz,
+        tz: TimeZones,
         axis: AxisIndex = ...,
         level: Level | None = ...,
         copy: _bool = ...,
     ) -> Series[S1]: ...
     def tz_localize(
         self,
-        tz,
+        tz: TimeZones,
         axis: AxisIndex = ...,
         level: Level | None = ...,
         copy: _bool = ...,

tests/test_scalars.py

@@ -1496,6 +1496,7 @@ def test_timestamp_misc_methods() -> None:
     check(assert_type(ts2, pd.Timestamp), pd.Timestamp)
     check(assert_type(ts.tz_localize("US/Pacific", False), pd.Timestamp), pd.Timestamp)
     check(assert_type(ts.tz_localize("US/Pacific", True), pd.Timestamp), pd.Timestamp)
+    check(assert_type(ts.tz_localize(1, True), pd.Timestamp), pd.Timestamp)
     check(assert_type(ts.tz_localize("US/Pacific", "NaT"), pd.Timestamp), pd.Timestamp)
     check(
         assert_type(ts.tz_localize("US/Pacific", "raise"), pd.Timestamp), pd.Timestamp

tests/test_timefuncs.py

@@ -468,6 +468,11 @@ def test_series_dt_accessors() -> None:
         pd.Series,
         pd.Timestamp,
     )
+    check(
+        assert_type(s0_local.dt.tz_convert(1), "TimestampSeries"),
+        pd.Series,
+        pd.Timestamp,
+    )
     check(
         assert_type(
             s0_local.dt.tz_convert(pytz.timezone("US/Eastern")), "TimestampSeries"
@@ -705,6 +710,19 @@ def test_some_offsets() -> None:
         assert_type(pd.date_range("1/1/2022", "2/1/2022", freq="1D"), pd.DatetimeIndex),
         pd.DatetimeIndex,
     )
+    check(
+        assert_type(
+            pd.date_range("1/1/2022", "2/1/2022", tz="Asia/Kathmandu", freq="1D"),
+            pd.DatetimeIndex,
+        ),
+        pd.DatetimeIndex,
+    )
+    check(
+        assert_type(
+            pd.date_range("1/1/2022", "2/1/2022", tz=3, freq="1D"), pd.DatetimeIndex
+        ),
+        pd.DatetimeIndex,
+    )
     check(
         assert_type(
             pd.date_range("1/1/2022", "2/1/2022", freq=Day()), pd.DatetimeIndex
@@ -717,6 +735,15 @@ def test_some_offsets() -> None:
         ),
         pd.DatetimeIndex,
     )
+    check(
+        assert_type(
+            pd.bdate_range(
+                "1/1/2022", "2/1/2022", tz="Asia/Kathmandu", freq=BusinessDay()
+            ),
+            pd.DatetimeIndex,
+        ),
+        pd.DatetimeIndex,
+    )
     # GH 755
     check(assert_type(dt.date.today() - Day(), pd.Timestamp), pd.Timestamp)
     check(assert_type(dt.date.today() + Day(), pd.Timestamp), pd.Timestamp)