From 33a7335b37180b210266c82653bc31a00240b9f9 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 13 Feb 2025 14:38:50 +0000 Subject: [PATCH] type `tz` in tz_convert, tz_localize, date_range, and bdate_range --- pandas-stubs/_libs/tslibs/timestamps.pyi | 7 +++--- pandas-stubs/_typing.pyi | 3 +++ pandas-stubs/core/arrays/datetimes.pyi | 7 ++++-- pandas-stubs/core/dtypes/dtypes.pyi | 5 ++--- pandas-stubs/core/frame.pyi | 5 +++-- pandas-stubs/core/generic.pyi | 10 --------- pandas-stubs/core/indexes/accessors.pyi | 5 +++-- pandas-stubs/core/indexes/datetimes.pyi | 7 +++--- pandas-stubs/core/series.pyi | 5 +++-- tests/test_scalars.py | 1 + tests/test_timefuncs.py | 27 ++++++++++++++++++++++++ 11 files changed, 55 insertions(+), 27 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index d99661b17..702c1287d 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/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: ... diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index c199a3718..322c1e622 100644 --- a/pandas-stubs/_typing.pyi +++ b/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"] diff --git a/pandas-stubs/core/arrays/datetimes.pyi b/pandas-stubs/core/arrays/datetimes.pyi index 7e4c7d7a1..f21af1d76 100644 --- a/pandas-stubs/core/arrays/datetimes.pyi +++ b/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=...): ... diff --git a/pandas-stubs/core/dtypes/dtypes.pyi b/pandas-stubs/core/dtypes/dtypes.pyi index 009134866..0abf8b2b9 100644 --- a/pandas-stubs/core/dtypes/dtypes.pyi +++ b/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 diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 540baf690..17f656043 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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 = ..., diff --git a/pandas-stubs/core/generic.pyi b/pandas-stubs/core/generic.pyi index 80eff9aa7..d9abd297a 100644 --- a/pandas-stubs/core/generic.pyi +++ b/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( diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index cc9d60b7f..1ce5dfdbc 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/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,7 +217,7 @@ 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"] @@ -224,7 +225,7 @@ class _DatetimeLikeNoTZMethods( | 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: ... diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index b0de6f5a3..5e89a8d8f 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/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 = ..., diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 70df87ab4..d2cc41c42 100644 --- a/pandas-stubs/core/series.pyi +++ b/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 = ..., diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 284995b4f..69a3dbf7d 100644 --- a/tests/test_scalars.py +++ b/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 diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 9cc22006f..e8d2e4773 100644 --- a/tests/test_timefuncs.py +++ b/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)