From 831d1846b1465c9e5e62fdd03aecc1661c301e9f Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Mon, 12 Feb 2024 13:49:05 +0900 Subject: [PATCH 1/7] Add unit property and as_unit method to NaTType --- pandas-stubs/_libs/tslibs/nattype.pyi | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/_libs/tslibs/nattype.pyi b/pandas-stubs/_libs/tslibs/nattype.pyi index 0fc2ba3d2..e8028d3fe 100644 --- a/pandas-stubs/_libs/tslibs/nattype.pyi +++ b/pandas-stubs/_libs/tslibs/nattype.pyi @@ -3,9 +3,13 @@ from datetime import ( timedelta, tzinfo as _tzinfo, ) +from typing import Literal import numpy as np -from typing_extensions import TypeAlias +from typing_extensions import ( + Self, + TypeAlias, +) from pandas._libs.tslibs.period import Period @@ -121,3 +125,8 @@ class NaTType: __le__: _NatComparison __gt__: _NatComparison __ge__: _NatComparison + @property + def unit(self) -> str: ... + def as_unit( + self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ... + ) -> Self: ... From eb26f35bd58b133948462facf26eac3fc60af0df Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Mon, 12 Feb 2024 13:33:43 +0900 Subject: [PATCH 2/7] Add unit property to Timestamp --- pandas-stubs/_libs/tslibs/timestamps.pyi | 2 ++ tests/test_scalars.py | 1 + 2 files changed, 3 insertions(+) diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index a8bf36b87..d389d0d39 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/pandas-stubs/_libs/tslibs/timestamps.pyi @@ -310,3 +310,5 @@ class Timestamp(datetime): def days_in_month(self) -> int: ... @property def daysinmonth(self) -> int: ... + @property + def unit(self) -> str: ... diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 478aab2ee..673af9719 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -1189,6 +1189,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) def test_timestamp_add_sub() -> None: From 31e8269a3efbef8847a3e4691294fa144a3d46ab Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Mon, 12 Feb 2024 03:20:42 +0900 Subject: [PATCH 3/7] Add as_unit method to Timestamp --- pandas-stubs/_libs/tslibs/timestamps.pyi | 3 +++ tests/test_scalars.py | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index d389d0d39..d6823d342 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/pandas-stubs/_libs/tslibs/timestamps.pyi @@ -312,3 +312,6 @@ class Timestamp(datetime): def daysinmonth(self) -> int: ... @property def unit(self) -> str: ... + def as_unit( + self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ... + ) -> Self: ... diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 673af9719..2a9f4348c 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -1646,6 +1646,11 @@ def test_timestamp_misc_methods() -> None: pd.Timestamp, ) + check(assert_type(ts2.as_unit("s"), pd.Timestamp), pd.Timestamp) + check(assert_type(ts2.as_unit("ms"), pd.Timestamp), pd.Timestamp) + check(assert_type(ts2.as_unit("us", round_ok=True), pd.Timestamp), pd.Timestamp) + check(assert_type(ts2.as_unit("ns", round_ok=False), pd.Timestamp), pd.Timestamp) + def test_timestamp_types_arithmetic() -> None: ts: pd.Timestamp = pd.to_datetime("2021-03-01") From a83501dc8b93e4119eaeb6629e10e4c9c06ed78f Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Mon, 12 Feb 2024 13:37:25 +0900 Subject: [PATCH 4/7] Add unit property to Timedelta --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 2 ++ tests/test_scalars.py | 1 + 2 files changed, 3 insertions(+) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 51dedb581..fb51a94ce 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -390,3 +390,5 @@ class Timedelta(timedelta): @property def components(self) -> Components: ... def view(self, dtype: npt.DTypeLike = ...) -> object: ... + @property + def unit(self) -> str: ... diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 2a9f4348c..18bae4efe 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -517,6 +517,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.ceil("D"), pd.Timedelta), pd.Timedelta) check(assert_type(td.floor(Day()), pd.Timedelta), pd.Timedelta) From 1dfb14c268eea186e17b7f6b99d752d88acdb681 Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Mon, 12 Feb 2024 13:48:00 +0900 Subject: [PATCH 5/7] Add as_unit method to Timedelta --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 3 +++ tests/test_scalars.py | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index fb51a94ce..65e48a686 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -392,3 +392,6 @@ class Timedelta(timedelta): 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: ... diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 18bae4efe..c71ef8e82 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -530,6 +530,11 @@ def test_timedelta_properties_methods() -> None: check(assert_type(td.view(np.int64), object), np.int64) check(assert_type(td.view("i8"), object), np.int64) + check(assert_type(td.as_unit("s"), pd.Timedelta), pd.Timedelta) + check(assert_type(td.as_unit("ms"), pd.Timedelta), pd.Timedelta) + check(assert_type(td.as_unit("us", round_ok=True), pd.Timedelta), pd.Timedelta) + check(assert_type(td.as_unit("ns", round_ok=False), pd.Timedelta), pd.Timedelta) + def test_timedelta_add_sub() -> None: td = pd.Timedelta("1 day") From b84dd6c886c43ae9a806d46da9de3658115b0a5b Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Tue, 13 Feb 2024 21:47:50 +0900 Subject: [PATCH 6/7] Add TimeUnit type alias --- pandas-stubs/_typing.pyi | 1 + 1 file changed, 1 insertion(+) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 89a7cbe72..9258054e3 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -775,6 +775,7 @@ RandomState: TypeAlias = ( | np.random.RandomState ) Frequency: TypeAlias = str | BaseOffset +TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] TimeGrouperOrigin: TypeAlias = ( Timestamp | Literal["epoch", "start", "start_day", "end", "end_day"] ) From 9fca549817ab51eede4ebecffc10e515357a007f Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Tue, 13 Feb 2024 22:13:43 +0900 Subject: [PATCH 7/7] Use TimeUnit type instead of str or Literal["s", "ms", "us", "ns"] --- pandas-stubs/_libs/tslibs/nattype.pyi | 8 +++----- pandas-stubs/_libs/tslibs/timedeltas.pyi | 11 ++++++----- pandas-stubs/_libs/tslibs/timestamps.pyi | 7 +++---- pandas-stubs/core/frame.pyi | 9 +++++---- pandas-stubs/core/indexes/datetimes.pyi | 8 +++----- pandas-stubs/core/series.pyi | 9 +++++---- pandas-stubs/io/json/_json.pyi | 9 +++++---- tests/test_scalars.py | 5 +++-- 8 files changed, 33 insertions(+), 33 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/nattype.pyi b/pandas-stubs/_libs/tslibs/nattype.pyi index e8028d3fe..cc74db157 100644 --- a/pandas-stubs/_libs/tslibs/nattype.pyi +++ b/pandas-stubs/_libs/tslibs/nattype.pyi @@ -3,7 +3,6 @@ from datetime import ( timedelta, tzinfo as _tzinfo, ) -from typing import Literal import numpy as np from typing_extensions import ( @@ -12,6 +11,7 @@ from typing_extensions import ( ) from pandas._libs.tslibs.period import Period +from pandas._typing import TimeUnit NaT: NaTType iNaT: int @@ -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: ... diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 65e48a686..13fc53b8d 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -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 @@ -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: ... diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index d6823d342..af856ae6f 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/pandas-stubs/_libs/tslibs/timestamps.pyi @@ -38,6 +38,7 @@ from pandas._libs.tslibs import ( Timedelta, ) from pandas._typing import ( + TimeUnit, np_ndarray_bool, npt, ) @@ -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: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index faae35994..d64d1cee7 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -115,6 +115,7 @@ from pandas._typing import ( StrLike, Suffixes, TimestampConvention, + TimeUnit, ValidationOptions, WriteBuffer, XMLParsers, @@ -2082,7 +2083,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 ) = ..., @@ -2101,7 +2102,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 ) = ..., @@ -2119,7 +2120,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 ) = ..., @@ -2137,7 +2138,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 ) = ..., diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index 50c2778ad..6d23fe56a 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/pandas-stubs/core/indexes/datetimes.pyi @@ -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 ( @@ -32,6 +29,7 @@ from pandas._typing import ( ArrayLike, DateAndDatetimeLike, IntervalClosedType, + TimeUnit, ) from pandas.core.dtypes.dtypes import DatetimeTZDtype @@ -100,7 +98,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( diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 383c8df99..45fb2b20e 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -144,6 +144,7 @@ from pandas._typing import ( TimedeltaDtypeArg, TimestampConvention, TimestampDtypeArg, + TimeUnit, UIntDtypeArg, VoidDtypeArg, WriteBuffer, @@ -486,7 +487,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 ) = ..., @@ -505,7 +506,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 ) = ..., @@ -523,7 +524,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 ) = ..., @@ -541,7 +542,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 ) = ..., diff --git a/pandas-stubs/io/json/_json.pyi b/pandas-stubs/io/json/_json.pyi index ea1fcb5bd..8151a6459 100644 --- a/pandas-stubs/io/json/_json.pyi +++ b/pandas-stubs/io/json/_json.pyi @@ -22,6 +22,7 @@ from pandas._typing import ( NDFrameT, ReadBuffer, StorageOptions, + TimeUnit, ) @overload @@ -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"] @@ -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"] @@ -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"] @@ -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"] diff --git a/tests/test_scalars.py b/tests/test_scalars.py index c71ef8e82..dc6e0e5e4 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -24,6 +24,7 @@ NaTType, ) from pandas._libs.tslibs.timedeltas import Components +from pandas._typing import TimeUnit from tests import ( TYPE_CHECKING_INVALID_USAGE, @@ -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) @@ -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: