Skip to content

Commit 9fca549

Browse files
committed
Use TimeUnit type instead of str or Literal["s", "ms", "us", "ns"]
1 parent b84dd6c commit 9fca549

File tree

8 files changed

+33
-33
lines changed

8 files changed

+33
-33
lines changed

pandas-stubs/_libs/tslibs/nattype.pyi

+3-5
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ from datetime import (
33
timedelta,
44
tzinfo as _tzinfo,
55
)
6-
from typing import Literal
76

87
import numpy as np
98
from typing_extensions import (
@@ -12,6 +11,7 @@ from typing_extensions import (
1211
)
1312

1413
from pandas._libs.tslibs.period import Period
14+
from pandas._typing import TimeUnit
1515

1616
NaT: NaTType
1717
iNaT: int
@@ -126,7 +126,5 @@ class NaTType:
126126
__gt__: _NatComparison
127127
__ge__: _NatComparison
128128
@property
129-
def unit(self) -> str: ...
130-
def as_unit(
131-
self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ...
132-
) -> Self: ...
129+
def unit(self) -> TimeUnit: ...
130+
def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...

pandas-stubs/_libs/tslibs/timedeltas.pyi

+6-5
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,10 @@ from pandas._libs.tslibs import (
3131
)
3232
from pandas._libs.tslibs.period import Period
3333
from pandas._libs.tslibs.timestamps import Timestamp
34-
from pandas._typing import npt
34+
from pandas._typing import (
35+
TimeUnit,
36+
npt,
37+
)
3538

3639
class Components(NamedTuple):
3740
days: int
@@ -391,7 +394,5 @@ class Timedelta(timedelta):
391394
def components(self) -> Components: ...
392395
def view(self, dtype: npt.DTypeLike = ...) -> object: ...
393396
@property
394-
def unit(self) -> str: ...
395-
def as_unit(
396-
self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ...
397-
) -> Self: ...
397+
def unit(self) -> TimeUnit: ...
398+
def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...

pandas-stubs/_libs/tslibs/timestamps.pyi

+3-4
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ from pandas._libs.tslibs import (
3838
Timedelta,
3939
)
4040
from pandas._typing import (
41+
TimeUnit,
4142
np_ndarray_bool,
4243
npt,
4344
)
@@ -311,7 +312,5 @@ class Timestamp(datetime):
311312
@property
312313
def daysinmonth(self) -> int: ...
313314
@property
314-
def unit(self) -> str: ...
315-
def as_unit(
316-
self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ...
317-
) -> Self: ...
315+
def unit(self) -> TimeUnit: ...
316+
def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...

pandas-stubs/core/frame.pyi

+5-4
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,7 @@ from pandas._typing import (
115115
StrLike,
116116
Suffixes,
117117
TimestampConvention,
118+
TimeUnit,
118119
ValidationOptions,
119120
WriteBuffer,
120121
XMLParsers,
@@ -2082,7 +2083,7 @@ class DataFrame(NDFrame, OpsMixin):
20822083
date_format: Literal["epoch", "iso"] | None = ...,
20832084
double_precision: int = ...,
20842085
force_ascii: _bool = ...,
2085-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
2086+
date_unit: TimeUnit = ...,
20862087
default_handler: (
20872088
Callable[[Any], _str | float | _bool | list | dict] | None
20882089
) = ...,
@@ -2101,7 +2102,7 @@ class DataFrame(NDFrame, OpsMixin):
21012102
date_format: Literal["epoch", "iso"] | None = ...,
21022103
double_precision: int = ...,
21032104
force_ascii: _bool = ...,
2104-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
2105+
date_unit: TimeUnit = ...,
21052106
default_handler: (
21062107
Callable[[Any], _str | float | _bool | list | dict] | None
21072108
) = ...,
@@ -2119,7 +2120,7 @@ class DataFrame(NDFrame, OpsMixin):
21192120
date_format: Literal["epoch", "iso"] | None = ...,
21202121
double_precision: int = ...,
21212122
force_ascii: _bool = ...,
2122-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
2123+
date_unit: TimeUnit = ...,
21232124
default_handler: (
21242125
Callable[[Any], _str | float | _bool | list | dict] | None
21252126
) = ...,
@@ -2137,7 +2138,7 @@ class DataFrame(NDFrame, OpsMixin):
21372138
date_format: Literal["epoch", "iso"] | None = ...,
21382139
double_precision: int = ...,
21392140
force_ascii: _bool = ...,
2140-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
2141+
date_unit: TimeUnit = ...,
21412142
default_handler: (
21422143
Callable[[Any], _str | float | _bool | list | dict] | None
21432144
) = ...,

pandas-stubs/core/indexes/datetimes.pyi

+3-5
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,7 @@ from datetime import (
77
timedelta,
88
tzinfo,
99
)
10-
from typing import (
11-
Literal,
12-
overload,
13-
)
10+
from typing import overload
1411

1512
import numpy as np
1613
from pandas import (
@@ -32,6 +29,7 @@ from pandas._typing import (
3229
ArrayLike,
3330
DateAndDatetimeLike,
3431
IntervalClosedType,
32+
TimeUnit,
3533
)
3634

3735
from pandas.core.dtypes.dtypes import DatetimeTZDtype
@@ -100,7 +98,7 @@ def date_range(
10098
normalize: bool = ...,
10199
name: Hashable | None = ...,
102100
inclusive: IntervalClosedType = ...,
103-
unit: Literal["s", "ms", "us", "ns"] | None = ...,
101+
unit: TimeUnit | None = ...,
104102
) -> DatetimeIndex: ...
105103
@overload
106104
def bdate_range(

pandas-stubs/core/series.pyi

+5-4
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ from pandas._typing import (
144144
TimedeltaDtypeArg,
145145
TimestampConvention,
146146
TimestampDtypeArg,
147+
TimeUnit,
147148
UIntDtypeArg,
148149
VoidDtypeArg,
149150
WriteBuffer,
@@ -486,7 +487,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
486487
date_format: Literal["epoch", "iso"] | None = ...,
487488
double_precision: int = ...,
488489
force_ascii: _bool = ...,
489-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
490+
date_unit: TimeUnit = ...,
490491
default_handler: (
491492
Callable[[Any], _str | float | _bool | list | dict] | None
492493
) = ...,
@@ -505,7 +506,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
505506
date_format: Literal["epoch", "iso"] | None = ...,
506507
double_precision: int = ...,
507508
force_ascii: _bool = ...,
508-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
509+
date_unit: TimeUnit = ...,
509510
default_handler: (
510511
Callable[[Any], _str | float | _bool | list | dict] | None
511512
) = ...,
@@ -523,7 +524,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
523524
date_format: Literal["epoch", "iso"] | None = ...,
524525
double_precision: int = ...,
525526
force_ascii: _bool = ...,
526-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
527+
date_unit: TimeUnit = ...,
527528
default_handler: (
528529
Callable[[Any], _str | float | _bool | list | dict] | None
529530
) = ...,
@@ -541,7 +542,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
541542
date_format: Literal["epoch", "iso"] | None = ...,
542543
double_precision: int = ...,
543544
force_ascii: _bool = ...,
544-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
545+
date_unit: TimeUnit = ...,
545546
default_handler: (
546547
Callable[[Any], _str | float | _bool | list | dict] | None
547548
) = ...,

pandas-stubs/io/json/_json.pyi

+5-4
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ from pandas._typing import (
2222
NDFrameT,
2323
ReadBuffer,
2424
StorageOptions,
25+
TimeUnit,
2526
)
2627

2728
@overload
@@ -35,7 +36,7 @@ def read_json(
3536
convert_dates: bool | list[str] = ...,
3637
keep_default_dates: bool = ...,
3738
precise_float: bool = ...,
38-
date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
39+
date_unit: TimeUnit | None = ...,
3940
encoding: str | None = ...,
4041
encoding_errors: (
4142
Literal["strict", "ignore", "replace", "backslashreplace", "surrogateescape"]
@@ -59,7 +60,7 @@ def read_json(
5960
convert_dates: bool | list[str] = ...,
6061
keep_default_dates: bool = ...,
6162
precise_float: bool = ...,
62-
date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
63+
date_unit: TimeUnit | None = ...,
6364
encoding: str | None = ...,
6465
encoding_errors: (
6566
Literal["strict", "ignore", "replace", "backslashreplace", "surrogateescape"]
@@ -83,7 +84,7 @@ def read_json(
8384
convert_dates: bool | list[str] = ...,
8485
keep_default_dates: bool = ...,
8586
precise_float: bool = ...,
86-
date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
87+
date_unit: TimeUnit | None = ...,
8788
encoding: str | None = ...,
8889
encoding_errors: (
8990
Literal["strict", "ignore", "replace", "backslashreplace", "surrogateescape"]
@@ -107,7 +108,7 @@ def read_json(
107108
convert_dates: bool | list[str] = ...,
108109
keep_default_dates: bool = ...,
109110
precise_float: bool = ...,
110-
date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
111+
date_unit: TimeUnit | None = ...,
111112
encoding: str | None = ...,
112113
encoding_errors: (
113114
Literal["strict", "ignore", "replace", "backslashreplace", "surrogateescape"]

tests/test_scalars.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
NaTType,
2525
)
2626
from pandas._libs.tslibs.timedeltas import Components
27+
from pandas._typing import TimeUnit
2728

2829
from tests import (
2930
TYPE_CHECKING_INVALID_USAGE,
@@ -517,7 +518,7 @@ def test_timedelta_properties_methods() -> None:
517518
check(assert_type(td.value, int), int)
518519
check(assert_type(td.resolution_string, str), str)
519520
check(assert_type(td.components, Components), Components)
520-
check(assert_type(td.unit, str), str)
521+
check(assert_type(td.unit, TimeUnit), str)
521522

522523
check(assert_type(td.ceil("D"), pd.Timedelta), pd.Timedelta)
523524
check(assert_type(td.floor(Day()), pd.Timedelta), pd.Timedelta)
@@ -1195,7 +1196,7 @@ def test_timestamp_properties() -> None:
11951196
check(assert_type(ts.tzinfo, Optional[dt.tzinfo]), type(None))
11961197
check(assert_type(ts.value, int), int)
11971198
check(assert_type(ts.year, int), int)
1198-
check(assert_type(ts.unit, str), str)
1199+
check(assert_type(ts.unit, TimeUnit), str)
11991200

12001201

12011202
def test_timestamp_add_sub() -> None:

0 commit comments

Comments
 (0)