Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 13, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions pandas-stubs/_libs/tslibs/timestamps.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ from typing import (
overload,
)

from _typing import TimeZones
import numpy as np
from pandas import (
DatetimeIndex,
Expand Down Expand Up @@ -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: ...
Expand Down Expand Up @@ -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: ...
Expand Down
3 changes: 3 additions & 0 deletions pandas-stubs/_typing.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ from collections.abc import (
Sequence,
)
import datetime
from datetime import tzinfo
from os import PathLike
import sys
from typing import (
Expand Down Expand Up @@ -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
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the related (private, for some reason?) alias in pandas is _TimeZones:

https://github.com/pandas-dev/pandas/blob/67872451d9a195c801c2395c92dab27edb3ecbbe/pandas/_libs/tslibs/timestamps.pyi#L30

I think TimeZone would be more accurate? I don't really mind, just explaining why I used the plural and where it comes from

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's fine


__all__ = ["npt", "type_t"]
7 changes: 5 additions & 2 deletions pandas-stubs/core/arrays/datetimes.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from datetime import tzinfo

from _typing import TimeZones
import numpy as np
from pandas.core.arrays.datetimelike import (
DatelikeOps,
Expand Down Expand Up @@ -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=...): ...
Expand Down
5 changes: 2 additions & 3 deletions pandas-stubs/core/dtypes/dtypes.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
5 changes: 3 additions & 2 deletions pandas-stubs/core/frame.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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 = ...,
Expand Down
10 changes: 0 additions & 10 deletions pandas-stubs/core/generic.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
5 changes: 3 additions & 2 deletions pandas-stubs/core/indexes/accessors.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ from typing import (
TypeVar,
)

from _typing import TimeZones
import numpy as np
import numpy.typing as npt
from pandas import (
Expand Down Expand Up @@ -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: ...
Expand Down
7 changes: 4 additions & 3 deletions pandas-stubs/core/indexes/datetimes.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ from datetime import (
)
from typing import overload

from _typing import TimeZones
import numpy as np
from pandas import (
DataFrame,
Expand Down Expand Up @@ -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 = ...,
Expand All @@ -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 = ...,
Expand All @@ -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 = ...,
Expand Down
5 changes: 3 additions & 2 deletions pandas-stubs/core/series.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ from typing import (
overload,
)

from _typing import TimeZones
from matplotlib.axes import (
Axes as PlotAxes,
SubplotBase,
Expand Down Expand Up @@ -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 = ...,
Expand Down
1 change: 1 addition & 0 deletions tests/test_scalars.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
27 changes: 27 additions & 0 deletions tests/test_timefuncs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down