Skip to content

Commit 33a7335

Browse files
MarcoGorelliMarcoGorelli
committed
type tz in tz_convert, tz_localize, date_range, and bdate_range
1 parent 4976e11 commit 33a7335

File tree

11 files changed

+55
-27
lines changed

11 files changed

+55
-27
lines changed

‎pandas-stubs/_libs/tslibs/timestamps.pyi‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ from typing import (
1515
overload,
1616
)
1717

18+
from _typing import TimeZones
1819
import numpy as np
1920
from pandas import (
2021
DatetimeIndex,
@@ -70,7 +71,7 @@ class Timestamp(datetime, SupportsIndex):
7071
tzinfo: _tzinfo | None = ...,
7172
*,
7273
nanosecond: int | None = ...,
73-
tz: str | _tzinfo | int | None = ...,
74+
tz: TimeZones = ...,
7475
unit: str | int | None = ...,
7576
fold: Literal[0, 1] | None = ...,
7677
) -> Self: ...
@@ -258,11 +259,11 @@ class Timestamp(datetime, SupportsIndex):
258259
def to_julian_date(self) -> np.float64: ...
259260
@property
260261
def asm8(self) -> np.datetime64: ...
261-
def tz_convert(self, tz: _tzinfo | str | None) -> Self: ...
262+
def tz_convert(self, tz: TimeZones) -> Self: ...
262263
# TODO: could return NaT?
263264
def tz_localize(
264265
self,
265-
tz: _tzinfo | str | None,
266+
tz: TimeZones,
266267
ambiguous: _Ambiguous = ...,
267268
nonexistent: _Nonexistent = ...,
268269
) -> Self: ...

‎pandas-stubs/_typing.pyi‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ from collections.abc import (
88
Sequence,
99
)
1010
import datetime
11+
from datetime import tzinfo
1112
from os import PathLike
1213
import sys
1314
from typing import (
@@ -820,4 +821,6 @@ TimeGrouperOrigin: TypeAlias = (
820821
ExcelReadEngine: TypeAlias = Literal["xlrd", "openpyxl", "odf", "pyxlsb", "calamine"]
821822
ExcelWriteEngine: TypeAlias = Literal["openpyxl", "odf", "xlsxwriter"]
822823

824+
TimeZones: TypeAlias = str | tzinfo | None | int
825+
823826
__all__ = ["npt", "type_t"]

‎pandas-stubs/core/arrays/datetimes.pyi‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from datetime import tzinfo
22

3+
from _typing import TimeZones
34
import numpy as np
45
from pandas.core.arrays.datetimelike import (
56
DatelikeOps,
@@ -28,8 +29,10 @@ class DatetimeArray(DatetimeLikeArrayMixin, TimelikeOps, DatelikeOps):
2829
def __array__(self, dtype=...) -> np.ndarray: ...
2930
def __iter__(self): ...
3031
def astype(self, dtype, copy: bool = ...): ...
31-
def tz_convert(self, tz): ...
32-
def tz_localize(self, tz, ambiguous: str = ..., nonexistent: str = ...): ...
32+
def tz_convert(self, tz: TimeZones): ...
33+
def tz_localize(
34+
self, tz: TimeZones, ambiguous: str = ..., nonexistent: str = ...
35+
): ...
3336
def to_pydatetime(self): ...
3437
def normalize(self): ...
3538
def to_period(self, freq=...): ...

‎pandas-stubs/core/dtypes/dtypes.pyi‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ from typing import (
55
TypeVar,
66
)
77

8+
from _typing import TimeZones
89
import numpy as np
910
from pandas.core.indexes.base import Index
1011
from pandas.core.series import Series
@@ -41,9 +42,7 @@ class CategoricalDtype(PandasExtensionDtype, ExtensionDtype):
4142
def ordered(self) -> Ordered: ...
4243

4344
class DatetimeTZDtype(PandasExtensionDtype):
44-
def __init__(
45-
self, unit: Literal["ns"] = ..., tz: str | int | dt.tzinfo | None = ...
46-
) -> None: ...
45+
def __init__(self, unit: Literal["ns"] = ..., tz: TimeZones = ...) -> None: ...
4746
@property
4847
def unit(self) -> Literal["ns"]: ...
4948
@property

‎pandas-stubs/core/frame.pyi‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ from typing import (
1919
overload,
2020
)
2121

22+
from _typing import TimeZones
2223
from matplotlib.axes import Axes as PlotAxes
2324
import numpy as np
2425
from pandas import (
@@ -2446,14 +2447,14 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
24462447
) -> Self: ...
24472448
def tz_convert(
24482449
self,
2449-
tz,
2450+
tz: TimeZones,
24502451
axis: Axis = ...,
24512452
level: Level | None = ...,
24522453
copy: _bool = ...,
24532454
) -> Self: ...
24542455
def tz_localize(
24552456
self,
2456-
tz,
2457+
tz: TimeZones,
24572458
axis: Axis = ...,
24582459
level: Level | None = ...,
24592460
copy: _bool = ...,

‎pandas-stubs/core/generic.pyi‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -546,16 +546,6 @@ class NDFrame(indexing.IndexingMixin):
546546
): ...
547547
def shift(self, periods=..., freq=..., axis=..., fill_value=...) -> Self: ...
548548
def truncate(self, before=..., after=..., axis=..., copy: _bool = ...) -> Self: ...
549-
def tz_convert(self, tz, axis=..., level=..., copy: _bool = ...) -> Self: ...
550-
def tz_localize(
551-
self,
552-
tz,
553-
axis=...,
554-
level=...,
555-
copy: _bool = ...,
556-
ambiguous=...,
557-
nonexistent: str = ...,
558-
) -> Self: ...
559549
def abs(self) -> Self: ...
560550
def describe(self, percentiles=..., include=..., exclude=...) -> NDFrame: ...
561551
def pct_change(

‎pandas-stubs/core/indexes/accessors.pyi‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ from typing import (
99
TypeVar,
1010
)
1111

12+
from _typing import TimeZones
1213
import numpy as np
1314
import numpy.typing as npt
1415
from pandas import (
@@ -216,15 +217,15 @@ class _DatetimeLikeNoTZMethods(
216217
) -> _DTToPeriodReturnType: ...
217218
def tz_localize(
218219
self,
219-
tz: tzinfo | str | None,
220+
tz: TimeZones,
220221
ambiguous: Literal["raise", "infer", "NaT"] | np_ndarray_bool = ...,
221222
nonexistent: (
222223
Literal["shift_forward", "shift_backward", "NaT", "raise"]
223224
| timedelta
224225
| Timedelta
225226
) = ...,
226227
) -> _DTNormalizeReturnType: ...
227-
def tz_convert(self, tz: tzinfo | str | None) -> _DTNormalizeReturnType: ...
228+
def tz_convert(self, tz: TimeZones) -> _DTNormalizeReturnType: ...
228229
def normalize(self) -> _DTNormalizeReturnType: ...
229230
def strftime(self, date_format: str) -> _DTStrKindReturnType: ...
230231
def month_name(self, locale: str | None = ...) -> _DTStrKindReturnType: ...

‎pandas-stubs/core/indexes/datetimes.pyi‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ from datetime import (
99
)
1010
from typing import overload
1111

12+
from _typing import TimeZones
1213
import numpy as np
1314
from pandas import (
1415
DataFrame,
@@ -92,7 +93,7 @@ def date_range(
9293
end: str | DateAndDatetimeLike | None = ...,
9394
periods: int | None = ...,
9495
freq: str | timedelta | Timedelta | BaseOffset = ...,
95-
tz: str | tzinfo = ...,
96+
tz: TimeZones = ...,
9697
normalize: bool = ...,
9798
name: Hashable | None = ...,
9899
inclusive: IntervalClosedType = ...,
@@ -104,7 +105,7 @@ def bdate_range(
104105
end: str | DateAndDatetimeLike | None = ...,
105106
periods: int | None = ...,
106107
freq: str | timedelta | Timedelta | BaseOffset = ...,
107-
tz: str | tzinfo = ...,
108+
tz: TimeZones = ...,
108109
normalize: bool = ...,
109110
name: Hashable | None = ...,
110111
weekmask: str | None = ...,
@@ -118,7 +119,7 @@ def bdate_range(
118119
periods: int | None = ...,
119120
*,
120121
freq: str | timedelta | Timedelta | BaseOffset,
121-
tz: str | tzinfo = ...,
122+
tz: TimeZones = ...,
122123
normalize: bool = ...,
123124
name: Hashable | None = ...,
124125
weekmask: str | None = ...,

‎pandas-stubs/core/series.pyi‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ from typing import (
2424
overload,
2525
)
2626

27+
from _typing import TimeZones
2728
from matplotlib.axes import (
2829
Axes as PlotAxes,
2930
SubplotBase,
@@ -1550,14 +1551,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
15501551
) -> Series[S1]: ...
15511552
def tz_convert(
15521553
self,
1553-
tz,
1554+
tz: TimeZones,
15541555
axis: AxisIndex = ...,
15551556
level: Level | None = ...,
15561557
copy: _bool = ...,
15571558
) -> Series[S1]: ...
15581559
def tz_localize(
15591560
self,
1560-
tz,
1561+
tz: TimeZones,
15611562
axis: AxisIndex = ...,
15621563
level: Level | None = ...,
15631564
copy: _bool = ...,

‎tests/test_scalars.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1496,6 +1496,7 @@ def test_timestamp_misc_methods() -> None:
14961496
check(assert_type(ts2, pd.Timestamp), pd.Timestamp)
14971497
check(assert_type(ts.tz_localize("US/Pacific", False), pd.Timestamp), pd.Timestamp)
14981498
check(assert_type(ts.tz_localize("US/Pacific", True), pd.Timestamp), pd.Timestamp)
1499+
check(assert_type(ts.tz_localize(1, True), pd.Timestamp), pd.Timestamp)
14991500
check(assert_type(ts.tz_localize("US/Pacific", "NaT"), pd.Timestamp), pd.Timestamp)
15001501
check(
15011502
assert_type(ts.tz_localize("US/Pacific", "raise"), pd.Timestamp), pd.Timestamp

0 commit comments

Comments
 (0)