Skip to content

Commit 3517aa2

Browse files
committed
Use TimeUnit type instead of str or Literal["s", "ms", "us", "ns"]
1 parent 0172f7a commit 3517aa2

File tree

8 files changed

+33
-33
lines changed

8 files changed

+33
-33
lines changed

pandas-stubs/_libs/tslibs/nattype.pyi

Lines changed: 3 additions & 5 deletions
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

Lines changed: 6 additions & 5 deletions
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

Lines changed: 3 additions & 4 deletions
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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ from pandas._typing import (
114114
StrLike,
115115
Suffixes,
116116
TimestampConvention,
117+
TimeUnit,
117118
ValidationOptions,
118119
WriteBuffer,
119120
XMLParsers,
@@ -2069,7 +2070,7 @@ class DataFrame(NDFrame, OpsMixin):
20692070
date_format: Literal["epoch", "iso"] | None = ...,
20702071
double_precision: int = ...,
20712072
force_ascii: _bool = ...,
2072-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
2073+
date_unit: TimeUnit = ...,
20732074
default_handler: Callable[[Any], _str | float | _bool | list | dict]
20742075
| None = ...,
20752076
lines: Literal[True],
@@ -2087,7 +2088,7 @@ class DataFrame(NDFrame, OpsMixin):
20872088
date_format: Literal["epoch", "iso"] | None = ...,
20882089
double_precision: int = ...,
20892090
force_ascii: _bool = ...,
2090-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
2091+
date_unit: TimeUnit = ...,
20912092
default_handler: Callable[[Any], _str | float | _bool | list | dict]
20922093
| None = ...,
20932094
lines: Literal[True],
@@ -2104,7 +2105,7 @@ class DataFrame(NDFrame, OpsMixin):
21042105
date_format: Literal["epoch", "iso"] | None = ...,
21052106
double_precision: int = ...,
21062107
force_ascii: _bool = ...,
2107-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
2108+
date_unit: TimeUnit = ...,
21082109
default_handler: Callable[[Any], _str | float | _bool | list | dict]
21092110
| None = ...,
21102111
lines: _bool = ...,
@@ -2121,7 +2122,7 @@ class DataFrame(NDFrame, OpsMixin):
21212122
date_format: Literal["epoch", "iso"] | None = ...,
21222123
double_precision: int = ...,
21232124
force_ascii: _bool = ...,
2124-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
2125+
date_unit: TimeUnit = ...,
21252126
default_handler: Callable[[Any], _str | float | _bool | list | dict]
21262127
| None = ...,
21272128
lines: _bool = ...,

pandas-stubs/core/indexes/datetimes.pyi

Lines changed: 3 additions & 5 deletions
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
@@ -102,7 +100,7 @@ def date_range(
102100
normalize: bool = ...,
103101
name: Hashable | None = ...,
104102
inclusive: IntervalClosedType = ...,
105-
unit: Literal["s", "ms", "us", "ns"] | None = ...,
103+
unit: TimeUnit | None = ...,
106104
) -> DatetimeIndex: ...
107105
@overload
108106
def bdate_range(

pandas-stubs/core/series.pyi

Lines changed: 5 additions & 4 deletions
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,
@@ -476,7 +477,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
476477
date_format: Literal["epoch", "iso"] | None = ...,
477478
double_precision: int = ...,
478479
force_ascii: _bool = ...,
479-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
480+
date_unit: TimeUnit = ...,
480481
default_handler: Callable[[Any], _str | float | _bool | list | dict]
481482
| None = ...,
482483
lines: Literal[True],
@@ -494,7 +495,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
494495
date_format: Literal["epoch", "iso"] | None = ...,
495496
double_precision: int = ...,
496497
force_ascii: _bool = ...,
497-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
498+
date_unit: TimeUnit = ...,
498499
default_handler: Callable[[Any], _str | float | _bool | list | dict]
499500
| None = ...,
500501
lines: Literal[True],
@@ -511,7 +512,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
511512
date_format: Literal["epoch", "iso"] | None = ...,
512513
double_precision: int = ...,
513514
force_ascii: _bool = ...,
514-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
515+
date_unit: TimeUnit = ...,
515516
default_handler: Callable[[Any], _str | float | _bool | list | dict]
516517
| None = ...,
517518
lines: _bool = ...,
@@ -528,7 +529,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
528529
date_format: Literal["epoch", "iso"] | None = ...,
529530
double_precision: int = ...,
530531
force_ascii: _bool = ...,
531-
date_unit: Literal["s", "ms", "us", "ns"] = ...,
532+
date_unit: TimeUnit = ...,
532533
default_handler: Callable[[Any], _str | float | _bool | list | dict]
533534
| None = ...,
534535
lines: _bool = ...,

pandas-stubs/io/json/_json.pyi

Lines changed: 5 additions & 4 deletions
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: Literal[
4142
"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: Literal[
6566
"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: Literal[
8990
"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: Literal[
113114
"strict", "ignore", "replace", "backslashreplace", "surrogateescape"

tests/test_scalars.py

Lines changed: 3 additions & 2 deletions
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)