From 5a37ab98e8af91304f2269be3722072d2c7d9319 Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Mon, 12 Feb 2024 03:38:03 +0900 Subject: [PATCH 1/4] Add unit property and as_unit method to DatetimeIndex and TimedeltaIndex --- pandas-stubs/core/indexes/datetimelike.pyi | 8 +++++++- tests/test_timefuncs.py | 10 ++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/indexes/datetimelike.pyi b/pandas-stubs/core/indexes/datetimelike.pyi index 5b138b744..2295140a8 100644 --- a/pandas-stubs/core/indexes/datetimelike.pyi +++ b/pandas-stubs/core/indexes/datetimelike.pyi @@ -1,5 +1,8 @@ +from typing import Literal + from pandas.core.indexes.extension import ExtensionIndex from pandas.core.indexes.timedeltas import TimedeltaIndex +from typing_extensions import Self from pandas._libs.tslibs import BaseOffset from pandas._typing import S1 @@ -19,4 +22,7 @@ class DatetimeIndexOpsMixin(ExtensionIndex[S1]): self, other: DatetimeIndexOpsMixin ) -> TimedeltaIndex: ... -class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]): ... +class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]): + @property + def unit(self) -> str: ... + def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> Self: ... diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index eeb177b7c..9a864ba0f 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -522,6 +522,11 @@ def test_datetimeindex_accessors() -> None: check(assert_type(i0.month_name(), pd.Index), pd.Index, str) check(assert_type(i0.day_name(), pd.Index), pd.Index, str) check(assert_type(i0.is_normalized, bool), bool) + check(assert_type(i0.unit, str), str) + check(assert_type(i0.as_unit("s"), pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(i0.as_unit("ms"), pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(i0.as_unit("us"), pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(i0.as_unit("ns"), pd.DatetimeIndex), pd.DatetimeIndex) def test_timedeltaindex_accessors() -> None: @@ -542,6 +547,11 @@ def test_timedeltaindex_accessors() -> None: assert_type(i0.floor("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta ) check(assert_type(i0.ceil("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta) + check(assert_type(i0.unit, str), str) + check(assert_type(i0.as_unit("s"), pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(i0.as_unit("ms"), pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(i0.as_unit("us"), pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(i0.as_unit("ns"), pd.TimedeltaIndex), pd.TimedeltaIndex) def test_periodindex_accessors() -> None: From f571b68a140066551387a746dc2aa5bbf89aaaf8 Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Mon, 12 Feb 2024 15:28:01 +0900 Subject: [PATCH 2/4] Add unit property and as_unit method to TimelikeOps --- pandas-stubs/core/arrays/datetimelike.pyi | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pandas-stubs/core/arrays/datetimelike.pyi b/pandas-stubs/core/arrays/datetimelike.pyi index b5b5dba76..1651d9877 100644 --- a/pandas-stubs/core/arrays/datetimelike.pyi +++ b/pandas-stubs/core/arrays/datetimelike.pyi @@ -1,10 +1,12 @@ from collections.abc import Sequence +from typing import Literal import numpy as np from pandas.core.arrays.base import ( ExtensionArray, ExtensionOpsMixin, ) +from typing_extensions import Self from pandas._libs import ( NaT as NaT, @@ -15,6 +17,9 @@ class DatelikeOps: def strftime(self, date_format): ... class TimelikeOps: + @property + def unit(self) -> str: ... + def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> Self: ... def round(self, freq, ambiguous: str = ..., nonexistent: str = ...): ... def floor(self, freq, ambiguous: str = ..., nonexistent: str = ...): ... def ceil(self, freq, ambiguous: str = ..., nonexistent: str = ...): ... From 2ca04fadf3c2be39378cd9cd114bb303b372e737 Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Tue, 13 Feb 2024 23:47:17 +0900 Subject: [PATCH 3/4] Add unit property and as_unit method to TimestampSeries and TimedeltaSeries --- pandas-stubs/core/indexes/accessors.pyi | 36 +++++++++++++++---------- tests/test_timefuncs.py | 35 ++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index a974fa6a9..ad3826e6c 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -154,16 +154,16 @@ class _DatetimeLikeOps( # type of the series, we don't know which kind of series was ...ed # in to the dt accessor -_DTRoundingMethodReturnType = TypeVar( - "_DTRoundingMethodReturnType", +_DTTimestampTimedeltaReturnType = TypeVar( + "_DTTimestampTimedeltaReturnType", Series, - TimedeltaSeries, TimestampSeries, + TimedeltaSeries, DatetimeIndex, TimedeltaIndex, ) -class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): +class _DatetimeRoundingMethods(Generic[_DTTimestampTimedeltaReturnType]): def round( self, freq: str | BaseOffset | None, @@ -173,7 +173,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): | timedelta | Timedelta ) = ..., - ) -> _DTRoundingMethodReturnType: ... + ) -> _DTTimestampTimedeltaReturnType: ... def floor( self, freq: str | BaseOffset | None, @@ -183,7 +183,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): | timedelta | Timedelta ) = ..., - ) -> _DTRoundingMethodReturnType: ... + ) -> _DTTimestampTimedeltaReturnType: ... def ceil( self, freq: str | BaseOffset | None, @@ -193,7 +193,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): | timedelta | Timedelta ) = ..., - ) -> _DTRoundingMethodReturnType: ... + ) -> _DTTimestampTimedeltaReturnType: ... _DTNormalizeReturnType = TypeVar( "_DTNormalizeReturnType", TimestampSeries, DatetimeIndex @@ -202,9 +202,9 @@ _DTStrKindReturnType = TypeVar("_DTStrKindReturnType", Series[str], Index) _DTToPeriodReturnType = TypeVar("_DTToPeriodReturnType", PeriodSeries, PeriodIndex) class _DatetimeLikeNoTZMethods( - _DatetimeRoundingMethods[_DTRoundingMethodReturnType], + _DatetimeRoundingMethods[_DTTimestampTimedeltaReturnType], Generic[ - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTNormalizeReturnType, _DTStrKindReturnType, _DTToPeriodReturnType, @@ -238,7 +238,7 @@ class _DatetimeNoTZProperties( _DTFreqReturnType, ], _DatetimeLikeNoTZMethods[ - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTNormalizeReturnType, _DTStrKindReturnType, _DTToPeriodReturnType, @@ -246,7 +246,7 @@ class _DatetimeNoTZProperties( Generic[ _DTFieldOpsReturnType, _DTBoolOpsReturnType, - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType, _DTFreqReturnType, @@ -261,7 +261,7 @@ class DatetimeProperties( _DatetimeNoTZProperties[ _DTFieldOpsReturnType, _DTBoolOpsReturnType, - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType, _DTFreqReturnType, @@ -272,7 +272,7 @@ class DatetimeProperties( Generic[ _DTFieldOpsReturnType, _DTBoolOpsReturnType, - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType, _DTFreqReturnType, @@ -283,6 +283,11 @@ class DatetimeProperties( ): def to_pydatetime(self) -> np.ndarray: ... def isocalendar(self) -> DataFrame: ... + @property + def unit(self) -> str: ... + def as_unit( + self, unit: Literal["s", "ms", "us", "ns"] + ) -> _DTTimestampTimedeltaReturnType: ... _TDNoRoundingMethodReturnType = TypeVar( "_TDNoRoundingMethodReturnType", Series[int], Index @@ -309,7 +314,10 @@ class TimedeltaProperties( Properties, _TimedeltaPropertiesNoRounding[Series[int], Series[float]], _DatetimeRoundingMethods[TimedeltaSeries], -): ... +): + @property + def unit(self) -> str: ... + def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> TimedeltaSeries: ... _PeriodDTReturnTypes = TypeVar("_PeriodDTReturnTypes", TimestampSeries, DatetimeIndex) _PeriodIntReturnTypes = TypeVar("_PeriodIntReturnTypes", Series[int], Index[int]) diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 9a864ba0f..e47545305 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -5,6 +5,7 @@ TYPE_CHECKING, Any, Optional, + cast, ) import numpy as np @@ -428,6 +429,11 @@ def test_series_dt_accessors() -> None: ) check(assert_type(s0.dt.month_name(), "pd.Series[str]"), pd.Series, str) check(assert_type(s0.dt.day_name(), "pd.Series[str]"), pd.Series, str) + check(assert_type(s0.dt.unit, str), str) + check(assert_type(s0.dt.as_unit("s"), "TimestampSeries"), pd.Series, pd.Timestamp) + check(assert_type(s0.dt.as_unit("ms"), "TimestampSeries"), pd.Series, pd.Timestamp) + check(assert_type(s0.dt.as_unit("us"), "TimestampSeries"), pd.Series, pd.Timestamp) + check(assert_type(s0.dt.as_unit("ns"), "TimestampSeries"), pd.Series, pd.Timestamp) i1 = pd.period_range(start="2022-06-01", periods=10) @@ -455,6 +461,35 @@ def test_series_dt_accessors() -> None: check(assert_type(s2.dt.components, pd.DataFrame), pd.DataFrame) check(assert_type(s2.dt.to_pytimedelta(), np.ndarray), np.ndarray) check(assert_type(s2.dt.total_seconds(), "pd.Series[float]"), pd.Series, float) + check(assert_type(s2.dt.unit, str), str) + check(assert_type(s2.dt.as_unit("s"), "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type(s2.dt.as_unit("ms"), "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type(s2.dt.as_unit("us"), "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type(s2.dt.as_unit("ns"), "TimedeltaSeries"), pd.Series, pd.Timedelta) + + # Checks for general Series other than TimestampSeries and TimedeltaSeries + + s4 = cast( + "pd.Series[pd.Timestamp]", + pd.Series([pd.Timestamp("2024-01-01"), pd.Timestamp("2024-01-02")]), + ) + + check(assert_type(s4.dt.unit, str), str) + check(assert_type(s4.dt.as_unit("s"), pd.Series), pd.Series, pd.Timestamp) + check(assert_type(s4.dt.as_unit("ms"), pd.Series), pd.Series, pd.Timestamp) + check(assert_type(s4.dt.as_unit("us"), pd.Series), pd.Series, pd.Timestamp) + check(assert_type(s4.dt.as_unit("ns"), pd.Series), pd.Series, pd.Timestamp) + + s5 = cast( + "pd.Series[pd.Timedelta]", + pd.Series([pd.Timedelta("1 day"), pd.Timedelta("2 days")]), + ) + + check(assert_type(s5.dt.unit, str), str) + check(assert_type(s5.dt.as_unit("s"), pd.Series), pd.Series, pd.Timedelta) + check(assert_type(s5.dt.as_unit("ms"), pd.Series), pd.Series, pd.Timedelta) + check(assert_type(s5.dt.as_unit("us"), pd.Series), pd.Series, pd.Timedelta) + check(assert_type(s5.dt.as_unit("ns"), pd.Series), pd.Series, pd.Timedelta) def test_datetimeindex_accessors() -> None: From 9120dce7249a100f5aee1ea0af72da834c86fc7c Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Tue, 13 Feb 2024 23:53:19 +0900 Subject: [PATCH 4/4] Use TimeUnit type alias instead of str or Literal["s", "ms", "us", "ns"] --- pandas-stubs/core/arrays/datetimelike.pyi | 6 +++--- pandas-stubs/core/indexes/accessors.pyi | 11 +++++------ pandas-stubs/core/indexes/datetimelike.pyi | 11 ++++++----- tests/test_timefuncs.py | 14 ++++++++------ 4 files changed, 22 insertions(+), 20 deletions(-) diff --git a/pandas-stubs/core/arrays/datetimelike.pyi b/pandas-stubs/core/arrays/datetimelike.pyi index 1651d9877..fc9548b0e 100644 --- a/pandas-stubs/core/arrays/datetimelike.pyi +++ b/pandas-stubs/core/arrays/datetimelike.pyi @@ -1,5 +1,4 @@ from collections.abc import Sequence -from typing import Literal import numpy as np from pandas.core.arrays.base import ( @@ -12,14 +11,15 @@ from pandas._libs import ( NaT as NaT, NaTType as NaTType, ) +from pandas._typing import TimeUnit class DatelikeOps: def strftime(self, date_format): ... class TimelikeOps: @property - def unit(self) -> str: ... - def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> Self: ... + def unit(self) -> TimeUnit: ... + def as_unit(self, unit: TimeUnit) -> Self: ... def round(self, freq, ambiguous: str = ..., nonexistent: str = ...): ... def floor(self, freq, ambiguous: str = ..., nonexistent: str = ...): ... def ceil(self, freq, ambiguous: str = ..., nonexistent: str = ...): ... diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index ad3826e6c..437b8b195 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -36,6 +36,7 @@ from pandas._libs.tslibs import BaseOffset from pandas._libs.tslibs.offsets import DateOffset from pandas._typing import ( TimestampConvention, + TimeUnit, np_ndarray_bool, ) @@ -284,10 +285,8 @@ class DatetimeProperties( def to_pydatetime(self) -> np.ndarray: ... def isocalendar(self) -> DataFrame: ... @property - def unit(self) -> str: ... - def as_unit( - self, unit: Literal["s", "ms", "us", "ns"] - ) -> _DTTimestampTimedeltaReturnType: ... + def unit(self) -> TimeUnit: ... + def as_unit(self, unit: TimeUnit) -> _DTTimestampTimedeltaReturnType: ... _TDNoRoundingMethodReturnType = TypeVar( "_TDNoRoundingMethodReturnType", Series[int], Index @@ -316,8 +315,8 @@ class TimedeltaProperties( _DatetimeRoundingMethods[TimedeltaSeries], ): @property - def unit(self) -> str: ... - def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> TimedeltaSeries: ... + def unit(self) -> TimeUnit: ... + def as_unit(self, unit: TimeUnit) -> TimedeltaSeries: ... _PeriodDTReturnTypes = TypeVar("_PeriodDTReturnTypes", TimestampSeries, DatetimeIndex) _PeriodIntReturnTypes = TypeVar("_PeriodIntReturnTypes", Series[int], Index[int]) diff --git a/pandas-stubs/core/indexes/datetimelike.pyi b/pandas-stubs/core/indexes/datetimelike.pyi index 2295140a8..fb27fae21 100644 --- a/pandas-stubs/core/indexes/datetimelike.pyi +++ b/pandas-stubs/core/indexes/datetimelike.pyi @@ -1,11 +1,12 @@ -from typing import Literal - from pandas.core.indexes.extension import ExtensionIndex from pandas.core.indexes.timedeltas import TimedeltaIndex from typing_extensions import Self from pandas._libs.tslibs import BaseOffset -from pandas._typing import S1 +from pandas._typing import ( + S1, + TimeUnit, +) class DatetimeIndexOpsMixin(ExtensionIndex[S1]): @property @@ -24,5 +25,5 @@ class DatetimeIndexOpsMixin(ExtensionIndex[S1]): class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]): @property - def unit(self) -> str: ... - def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> Self: ... + def unit(self) -> TimeUnit: ... + def as_unit(self, unit: TimeUnit) -> Self: ... diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index e47545305..36dd6d548 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -25,6 +25,8 @@ from pandas._typing import FulldatetimeDict else: FulldatetimeDict = Any +from pandas._typing import TimeUnit + from tests import ( TYPE_CHECKING_INVALID_USAGE, check, @@ -429,7 +431,7 @@ def test_series_dt_accessors() -> None: ) check(assert_type(s0.dt.month_name(), "pd.Series[str]"), pd.Series, str) check(assert_type(s0.dt.day_name(), "pd.Series[str]"), pd.Series, str) - check(assert_type(s0.dt.unit, str), str) + check(assert_type(s0.dt.unit, TimeUnit), str) check(assert_type(s0.dt.as_unit("s"), "TimestampSeries"), pd.Series, pd.Timestamp) check(assert_type(s0.dt.as_unit("ms"), "TimestampSeries"), pd.Series, pd.Timestamp) check(assert_type(s0.dt.as_unit("us"), "TimestampSeries"), pd.Series, pd.Timestamp) @@ -461,7 +463,7 @@ def test_series_dt_accessors() -> None: check(assert_type(s2.dt.components, pd.DataFrame), pd.DataFrame) check(assert_type(s2.dt.to_pytimedelta(), np.ndarray), np.ndarray) check(assert_type(s2.dt.total_seconds(), "pd.Series[float]"), pd.Series, float) - check(assert_type(s2.dt.unit, str), str) + check(assert_type(s2.dt.unit, TimeUnit), str) check(assert_type(s2.dt.as_unit("s"), "TimedeltaSeries"), pd.Series, pd.Timedelta) check(assert_type(s2.dt.as_unit("ms"), "TimedeltaSeries"), pd.Series, pd.Timedelta) check(assert_type(s2.dt.as_unit("us"), "TimedeltaSeries"), pd.Series, pd.Timedelta) @@ -474,7 +476,7 @@ def test_series_dt_accessors() -> None: pd.Series([pd.Timestamp("2024-01-01"), pd.Timestamp("2024-01-02")]), ) - check(assert_type(s4.dt.unit, str), str) + check(assert_type(s4.dt.unit, TimeUnit), str) check(assert_type(s4.dt.as_unit("s"), pd.Series), pd.Series, pd.Timestamp) check(assert_type(s4.dt.as_unit("ms"), pd.Series), pd.Series, pd.Timestamp) check(assert_type(s4.dt.as_unit("us"), pd.Series), pd.Series, pd.Timestamp) @@ -485,7 +487,7 @@ def test_series_dt_accessors() -> None: pd.Series([pd.Timedelta("1 day"), pd.Timedelta("2 days")]), ) - check(assert_type(s5.dt.unit, str), str) + check(assert_type(s5.dt.unit, TimeUnit), str) check(assert_type(s5.dt.as_unit("s"), pd.Series), pd.Series, pd.Timedelta) check(assert_type(s5.dt.as_unit("ms"), pd.Series), pd.Series, pd.Timedelta) check(assert_type(s5.dt.as_unit("us"), pd.Series), pd.Series, pd.Timedelta) @@ -557,7 +559,7 @@ def test_datetimeindex_accessors() -> None: check(assert_type(i0.month_name(), pd.Index), pd.Index, str) check(assert_type(i0.day_name(), pd.Index), pd.Index, str) check(assert_type(i0.is_normalized, bool), bool) - check(assert_type(i0.unit, str), str) + check(assert_type(i0.unit, TimeUnit), str) check(assert_type(i0.as_unit("s"), pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(i0.as_unit("ms"), pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(i0.as_unit("us"), pd.DatetimeIndex), pd.DatetimeIndex) @@ -582,7 +584,7 @@ def test_timedeltaindex_accessors() -> None: assert_type(i0.floor("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta ) check(assert_type(i0.ceil("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta) - check(assert_type(i0.unit, str), str) + check(assert_type(i0.unit, TimeUnit), str) check(assert_type(i0.as_unit("s"), pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(i0.as_unit("ms"), pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(i0.as_unit("us"), pd.TimedeltaIndex), pd.TimedeltaIndex)