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..823ed1d59 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1571,17 +1571,147 @@ def test_updated_astype() -> None: np.complex128, ) + 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()