From daac016e5be3f024dfeddfe3b95d6a5a7b52a847 Mon Sep 17 00:00:00 2001 From: Siddhartha Gandhi Date: Fri, 3 Mar 2023 13:44:15 -0500 Subject: [PATCH 1/3] Support all possible time units for Series.astype literals --- pandas-stubs/_typing.pyi | 35 +++++++++++++++++++++++++++++++++-- tests/test_series.py | 12 +++++++++++- 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index afa1a38a8..8af781e43 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -146,8 +146,39 @@ ComplexDtypeArg: TypeAlias = ( | type[np.complex128] | Literal["complex64", "complex128"] ) -TimedeltaDtypeArg: TypeAlias = Literal["timedelta64[ns]"] -TimestampDtypeArg: TypeAlias = Literal["datetime64[ns]"] +# Refer to https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units +TimedeltaDtypeArg: TypeAlias = Literal[ + "timedelta64[Y]", + "timedelta64[M]", + "timedelta64[W]", + "timedelta64[D]", + "timedelta64[h]", + "timedelta64[m]", + "timedelta64[s]", + "timedelta64[ms]", + "timedelta64[us]", + "timedelta64[μs]", + "timedelta64[ns]", + "timedelta64[ps]", + "timedelta64[fs]", + "timedelta64[as]", +] +TimestampDtypeArg: TypeAlias = Literal[ + "datetime64[Y]", + "datetime64[M]", + "datetime64[W]", + "datetime64[D]", + "datetime64[h]", + "datetime64[m]", + "datetime64[s]", + "datetime64[ms]", + "datetime64[us]", + "datetime64[μs]", + "datetime64[ns]", + "datetime64[ps]", + "datetime64[fs]", + "datetime64[as]", +] CategoryDtypeArg: TypeAlias = CategoricalDtype | Literal["category"] AstypeArg: TypeAlias = ( diff --git a/tests/test_series.py b/tests/test_series.py index 63c2cd011..9a8cbadd6 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1571,12 +1571,22 @@ def test_updated_astype() -> None: np.complex128, ) + # Check a couple of selected time units, but not all, to avoid excessive test cases + check( + assert_type(s.astype("timedelta64[M]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("datetime64[M]"), TimestampSeries), + pd.Series, + Timestamp, + ) check( assert_type(s.astype("timedelta64[ns]"), TimedeltaSeries), pd.Series, Timedelta, ) - check( assert_type(s.astype("datetime64[ns]"), TimestampSeries), pd.Series, From 8a9720b38e584409f0b1dbe9adf8bad942165888 Mon Sep 17 00:00:00 2001 From: Siddhartha Gandhi Date: Fri, 3 Mar 2023 16:16:02 -0500 Subject: [PATCH 2/3] Update test to check all timedelta/timestamp dtype args in a loop --- tests/test_series.py | 41 +++++++++++++++++++++-------------------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/tests/test_series.py b/tests/test_series.py index 9a8cbadd6..3e6abc2fc 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -16,6 +16,7 @@ Tuple, TypeVar, cast, + get_args, ) import numpy as np @@ -1572,26 +1573,26 @@ def test_updated_astype() -> None: ) # Check a couple of selected time units, but not all, to avoid excessive test cases - check( - assert_type(s.astype("timedelta64[M]"), TimedeltaSeries), - pd.Series, - Timedelta, - ) - check( - assert_type(s.astype("datetime64[M]"), TimestampSeries), - pd.Series, - Timestamp, - ) - check( - assert_type(s.astype("timedelta64[ns]"), TimedeltaSeries), - pd.Series, - Timedelta, - ) - check( - assert_type(s.astype("datetime64[ns]"), TimestampSeries), - pd.Series, - Timestamp, - ) + if TYPE_CHECKING: + from pandas._typing import ( + TimedeltaDtypeArg, + TimestampDtypeArg, + ) + + tdt_types = cast(Tuple[TimedeltaDtypeArg, ...], get_args(TimedeltaDtypeArg)) + for tdt_type in tdt_types: + check( + assert_type(s.astype(tdt_type), TimedeltaSeries), + pd.Series, + Timedelta, + ) + ts_types = cast(Tuple[TimestampDtypeArg, ...], get_args(TimestampDtypeArg)) + for ts_type in ts_types: + check( + assert_type(s.astype(ts_type), TimestampSeries), + pd.Series, + Timestamp, + ) orseries = pd.Series([Decimal(x) for x in [1, 2, 3]]) newtype = DecimalDtype() From 01a97f4b2fd33f2bce7b3c4421072356728c45f4 Mon Sep 17 00:00:00 2001 From: Siddhartha Gandhi Date: Fri, 3 Mar 2023 18:40:13 -0500 Subject: [PATCH 3/3] List out all of the individual string literals in astype test --- tests/test_series.py | 163 +++++++++++++++++++++++++++++++++++++------ 1 file changed, 141 insertions(+), 22 deletions(-) diff --git a/tests/test_series.py b/tests/test_series.py index 3e6abc2fc..823ed1d59 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -16,7 +16,6 @@ Tuple, TypeVar, cast, - get_args, ) import numpy as np @@ -1572,27 +1571,147 @@ def test_updated_astype() -> None: np.complex128, ) - # Check a couple of selected time units, but not all, to avoid excessive test cases - if TYPE_CHECKING: - from pandas._typing import ( - TimedeltaDtypeArg, - TimestampDtypeArg, - ) - - tdt_types = cast(Tuple[TimedeltaDtypeArg, ...], get_args(TimedeltaDtypeArg)) - for tdt_type in tdt_types: - check( - assert_type(s.astype(tdt_type), TimedeltaSeries), - pd.Series, - Timedelta, - ) - ts_types = cast(Tuple[TimestampDtypeArg, ...], get_args(TimestampDtypeArg)) - for ts_type in ts_types: - check( - assert_type(s.astype(ts_type), TimestampSeries), - pd.Series, - Timestamp, - ) + check( + assert_type(s.astype("timedelta64[Y]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[M]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[W]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[D]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[h]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[m]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[s]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[ms]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[us]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[μs]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[ns]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[ps]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[fs]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + check( + assert_type(s.astype("timedelta64[as]"), TimedeltaSeries), + pd.Series, + Timedelta, + ) + + check( + assert_type(s.astype("datetime64[Y]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[M]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[W]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[D]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[h]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[m]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[s]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[ms]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[us]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[μs]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[ns]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[ps]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[fs]"), TimestampSeries), + pd.Series, + Timestamp, + ) + check( + assert_type(s.astype("datetime64[as]"), TimestampSeries), + pd.Series, + Timestamp, + ) orseries = pd.Series([Decimal(x) for x in [1, 2, 3]]) newtype = DecimalDtype()