Skip to content

Commit

Permalink
Use TimeUnit type instead of str or Literal["s", "ms", "us", "ns"]
Browse files Browse the repository at this point in the history
  • Loading branch information
skatsuta committed Feb 13, 2024
1 parent 0172f7a commit 3517aa2
Show file tree
Hide file tree
Showing 8 changed files with 33 additions and 33 deletions.
8 changes: 3 additions & 5 deletions pandas-stubs/_libs/tslibs/nattype.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ from datetime import (
timedelta,
tzinfo as _tzinfo,
)
from typing import Literal

import numpy as np
from typing_extensions import (
Expand All @@ -12,6 +11,7 @@ from typing_extensions import (
)

from pandas._libs.tslibs.period import Period
from pandas._typing import TimeUnit

NaT: NaTType
iNaT: int
Expand Down Expand Up @@ -126,7 +126,5 @@ class NaTType:
__gt__: _NatComparison
__ge__: _NatComparison
@property
def unit(self) -> str: ...
def as_unit(
self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ...
) -> Self: ...
def unit(self) -> TimeUnit: ...
def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...
11 changes: 6 additions & 5 deletions pandas-stubs/_libs/tslibs/timedeltas.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,10 @@ from pandas._libs.tslibs import (
)
from pandas._libs.tslibs.period import Period
from pandas._libs.tslibs.timestamps import Timestamp
from pandas._typing import npt
from pandas._typing import (
TimeUnit,
npt,
)

class Components(NamedTuple):
days: int
Expand Down Expand Up @@ -391,7 +394,5 @@ class Timedelta(timedelta):
def components(self) -> Components: ...
def view(self, dtype: npt.DTypeLike = ...) -> object: ...
@property
def unit(self) -> str: ...
def as_unit(
self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ...
) -> Self: ...
def unit(self) -> TimeUnit: ...
def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...
7 changes: 3 additions & 4 deletions pandas-stubs/_libs/tslibs/timestamps.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ from pandas._libs.tslibs import (
Timedelta,
)
from pandas._typing import (
TimeUnit,
np_ndarray_bool,
npt,
)
Expand Down Expand Up @@ -311,7 +312,5 @@ class Timestamp(datetime):
@property
def daysinmonth(self) -> int: ...
@property
def unit(self) -> str: ...
def as_unit(
self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ...
) -> Self: ...
def unit(self) -> TimeUnit: ...
def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...
9 changes: 5 additions & 4 deletions pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,7 @@ from pandas._typing import (
StrLike,
Suffixes,
TimestampConvention,
TimeUnit,
ValidationOptions,
WriteBuffer,
XMLParsers,
Expand Down Expand Up @@ -2069,7 +2070,7 @@ class DataFrame(NDFrame, OpsMixin):
date_format: Literal["epoch", "iso"] | None = ...,
double_precision: int = ...,
force_ascii: _bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] = ...,
date_unit: TimeUnit = ...,
default_handler: Callable[[Any], _str | float | _bool | list | dict]
| None = ...,
lines: Literal[True],
Expand All @@ -2087,7 +2088,7 @@ class DataFrame(NDFrame, OpsMixin):
date_format: Literal["epoch", "iso"] | None = ...,
double_precision: int = ...,
force_ascii: _bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] = ...,
date_unit: TimeUnit = ...,
default_handler: Callable[[Any], _str | float | _bool | list | dict]
| None = ...,
lines: Literal[True],
Expand All @@ -2104,7 +2105,7 @@ class DataFrame(NDFrame, OpsMixin):
date_format: Literal["epoch", "iso"] | None = ...,
double_precision: int = ...,
force_ascii: _bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] = ...,
date_unit: TimeUnit = ...,
default_handler: Callable[[Any], _str | float | _bool | list | dict]
| None = ...,
lines: _bool = ...,
Expand All @@ -2121,7 +2122,7 @@ class DataFrame(NDFrame, OpsMixin):
date_format: Literal["epoch", "iso"] | None = ...,
double_precision: int = ...,
force_ascii: _bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] = ...,
date_unit: TimeUnit = ...,
default_handler: Callable[[Any], _str | float | _bool | list | dict]
| None = ...,
lines: _bool = ...,
Expand Down
8 changes: 3 additions & 5 deletions pandas-stubs/core/indexes/datetimes.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,7 @@ from datetime import (
timedelta,
tzinfo,
)
from typing import (
Literal,
overload,
)
from typing import overload

import numpy as np
from pandas import (
Expand All @@ -32,6 +29,7 @@ from pandas._typing import (
ArrayLike,
DateAndDatetimeLike,
IntervalClosedType,
TimeUnit,
)

from pandas.core.dtypes.dtypes import DatetimeTZDtype
Expand Down Expand Up @@ -102,7 +100,7 @@ def date_range(
normalize: bool = ...,
name: Hashable | None = ...,
inclusive: IntervalClosedType = ...,
unit: Literal["s", "ms", "us", "ns"] | None = ...,
unit: TimeUnit | None = ...,
) -> DatetimeIndex: ...
@overload
def bdate_range(
Expand Down
9 changes: 5 additions & 4 deletions pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,7 @@ from pandas._typing import (
TimedeltaDtypeArg,
TimestampConvention,
TimestampDtypeArg,
TimeUnit,
UIntDtypeArg,
VoidDtypeArg,
WriteBuffer,
Expand Down Expand Up @@ -476,7 +477,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
date_format: Literal["epoch", "iso"] | None = ...,
double_precision: int = ...,
force_ascii: _bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] = ...,
date_unit: TimeUnit = ...,
default_handler: Callable[[Any], _str | float | _bool | list | dict]
| None = ...,
lines: Literal[True],
Expand All @@ -494,7 +495,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
date_format: Literal["epoch", "iso"] | None = ...,
double_precision: int = ...,
force_ascii: _bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] = ...,
date_unit: TimeUnit = ...,
default_handler: Callable[[Any], _str | float | _bool | list | dict]
| None = ...,
lines: Literal[True],
Expand All @@ -511,7 +512,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
date_format: Literal["epoch", "iso"] | None = ...,
double_precision: int = ...,
force_ascii: _bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] = ...,
date_unit: TimeUnit = ...,
default_handler: Callable[[Any], _str | float | _bool | list | dict]
| None = ...,
lines: _bool = ...,
Expand All @@ -528,7 +529,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
date_format: Literal["epoch", "iso"] | None = ...,
double_precision: int = ...,
force_ascii: _bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] = ...,
date_unit: TimeUnit = ...,
default_handler: Callable[[Any], _str | float | _bool | list | dict]
| None = ...,
lines: _bool = ...,
Expand Down
9 changes: 5 additions & 4 deletions pandas-stubs/io/json/_json.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ from pandas._typing import (
NDFrameT,
ReadBuffer,
StorageOptions,
TimeUnit,
)

@overload
Expand All @@ -35,7 +36,7 @@ def read_json(
convert_dates: bool | list[str] = ...,
keep_default_dates: bool = ...,
precise_float: bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
date_unit: TimeUnit | None = ...,
encoding: str | None = ...,
encoding_errors: Literal[
"strict", "ignore", "replace", "backslashreplace", "surrogateescape"
Expand All @@ -59,7 +60,7 @@ def read_json(
convert_dates: bool | list[str] = ...,
keep_default_dates: bool = ...,
precise_float: bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
date_unit: TimeUnit | None = ...,
encoding: str | None = ...,
encoding_errors: Literal[
"strict", "ignore", "replace", "backslashreplace", "surrogateescape"
Expand All @@ -83,7 +84,7 @@ def read_json(
convert_dates: bool | list[str] = ...,
keep_default_dates: bool = ...,
precise_float: bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
date_unit: TimeUnit | None = ...,
encoding: str | None = ...,
encoding_errors: Literal[
"strict", "ignore", "replace", "backslashreplace", "surrogateescape"
Expand All @@ -107,7 +108,7 @@ def read_json(
convert_dates: bool | list[str] = ...,
keep_default_dates: bool = ...,
precise_float: bool = ...,
date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
date_unit: TimeUnit | None = ...,
encoding: str | None = ...,
encoding_errors: Literal[
"strict", "ignore", "replace", "backslashreplace", "surrogateescape"
Expand Down
5 changes: 3 additions & 2 deletions tests/test_scalars.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
NaTType,
)
from pandas._libs.tslibs.timedeltas import Components
from pandas._typing import TimeUnit

from tests import (
TYPE_CHECKING_INVALID_USAGE,
Expand Down Expand Up @@ -517,7 +518,7 @@ def test_timedelta_properties_methods() -> None:
check(assert_type(td.value, int), int)
check(assert_type(td.resolution_string, str), str)
check(assert_type(td.components, Components), Components)
check(assert_type(td.unit, str), str)
check(assert_type(td.unit, TimeUnit), str)

check(assert_type(td.ceil("D"), pd.Timedelta), pd.Timedelta)
check(assert_type(td.floor(Day()), pd.Timedelta), pd.Timedelta)
Expand Down Expand Up @@ -1195,7 +1196,7 @@ def test_timestamp_properties() -> None:
check(assert_type(ts.tzinfo, Optional[dt.tzinfo]), type(None))
check(assert_type(ts.value, int), int)
check(assert_type(ts.year, int), int)
check(assert_type(ts.unit, str), str)
check(assert_type(ts.unit, TimeUnit), str)


def test_timestamp_add_sub() -> None:
Expand Down

0 comments on commit 3517aa2

Please sign in to comment.