diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 6cb6bfa92..5ba453ecd 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -863,6 +863,10 @@ np_ndarray_num: TypeAlias = np_ndarray[ShapeT, NpNumT] # Numpy arrays with known shape (Do not use as argument types, only as return types) np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[GenericT]] +np_1darray_str: TypeAlias = np_1darray[np.str_] +np_1darray_bytes: TypeAlias = np_1darray[np.bytes_] +np_1darray_complex: TypeAlias = np_1darray[np.complexfloating] +np_1darray_object: TypeAlias = np_1darray[np.object_] np_1darray_bool: TypeAlias = np_1darray[np.bool] np_1darray_intp: TypeAlias = np_1darray[np.intp] np_1darray_int64: TypeAlias = np_1darray[np.int64] @@ -870,7 +874,6 @@ np_1darray_anyint: TypeAlias = np_1darray[np.integer] np_1darray_float: TypeAlias = np_1darray[np.floating] np_1darray_dt: TypeAlias = np_1darray[np.datetime64] np_1darray_td: TypeAlias = np_1darray[np.timedelta64] - np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]] DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic]) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 61eb0222e..96f73a950 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -220,10 +220,16 @@ from pandas._typing import ( WriteBuffer, _T_co, np_1darray, + np_1darray_anyint, np_1darray_bool, + np_1darray_bytes, + np_1darray_complex, np_1darray_dt, + np_1darray_float, np_1darray_int64, np_1darray_intp, + np_1darray_object, + np_1darray_str, np_1darray_td, np_ndarray, np_ndarray_anyint, @@ -4469,7 +4475,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = False, na_value: Scalar = ..., **kwargs: Any, - ) -> np_1darray[np.object_]: ... + ) -> np_1darray_object: ... @overload def to_numpy( self: Series[Period], @@ -4485,7 +4491,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = False, na_value: Scalar = ..., **kwargs: Any, - ) -> np_1darray[np.object_]: ... + ) -> np_1darray_object: ... @overload def to_numpy( self: Series[BaseOffset], @@ -4493,7 +4499,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = False, na_value: Scalar = ..., **kwargs: Any, - ) -> np_1darray[np.bytes_]: ... + ) -> np_1darray_bytes: ... @overload def to_numpy( self: Series[Interval], @@ -4501,7 +4507,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = False, na_value: Scalar = ..., **kwargs: Any, - ) -> np_1darray[np.object_]: ... + ) -> np_1darray_object: ... @overload def to_numpy( self: Series[Interval], @@ -4511,6 +4517,54 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): **kwargs: Any, ) -> np_1darray[_T_INTERVAL_NP]: ... @overload + def to_numpy( + self: Series[int], + dtype: DTypeLike | None = None, + copy: bool = False, + na_value: Scalar = ..., + **kwargs: Any, + ) -> np_1darray_anyint: ... + @overload + def to_numpy( + self: Series[float], + dtype: DTypeLike | None = None, + copy: bool = False, + na_value: Scalar = ..., + **kwargs: Any, + ) -> np_1darray_float: ... + @overload + def to_numpy( + self: Series[complex], + dtype: DTypeLike | None = None, + copy: bool = False, + na_value: Scalar = ..., + **kwargs: Any, + ) -> np_1darray_complex: ... + @overload + def to_numpy( + self: Series[bool], + dtype: DTypeLike | None = None, + copy: bool = False, + na_value: Scalar = ..., + **kwargs: Any, + ) -> np_1darray_bool: ... + @overload + def to_numpy( + self: Series[_str], + dtype: DTypeLike | None = None, + copy: bool = False, + na_value: Scalar = ..., + **kwargs: Any, + ) -> np_1darray_str: ... + @overload + def to_numpy( + self: Series[bytes], + dtype: DTypeLike | None = None, + copy: bool = False, + na_value: Scalar = ..., + **kwargs: Any, + ) -> np_1darray_bytes: ... + @overload def to_numpy( # pyright: ignore[reportIncompatibleMethodOverride] self, dtype: DTypeLike | None = None, diff --git a/tests/__init__.py b/tests/__init__.py index aa640305e..701e1757d 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -47,10 +47,14 @@ np_1darray as np_1darray, np_1darray_anyint as np_1darray_anyint, np_1darray_bool as np_1darray_bool, + np_1darray_bytes as np_1darray_bytes, + np_1darray_complex as np_1darray_complex, np_1darray_dt as np_1darray_dt, np_1darray_float as np_1darray_float, np_1darray_int64 as np_1darray_int64, np_1darray_intp as np_1darray_intp, + np_1darray_object as np_1darray_object, + np_1darray_str as np_1darray_str, np_1darray_td as np_1darray_td, np_2darray as np_2darray, np_ndarray as np_ndarray, @@ -67,7 +71,11 @@ _S = TypeVar("_S", bound=tuple[int, ...]) # Separately define here so pytest works np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[_G]] - np_1darray_bool: TypeAlias = np_1darray[np.bool] + np_1darray_bool: TypeAlias = np.ndarray[tuple[int], np.bool_] + np_1darray_str: TypeAlias = np.ndarray[tuple[int], np.str_] + np_1darray_bytes: TypeAlias = np.ndarray[tuple[int], np.bytes_] + np_1darray_complex: TypeAlias = np.ndarray[tuple[int], np.complexfloating] + np_1darray_object: TypeAlias = np.ndarray[tuple[int], np.object_] np_1darray_intp: TypeAlias = np_1darray[np.intp] np_1darray_int64: TypeAlias = np_1darray[np.int64] np_1darray_anyint: TypeAlias = np_1darray[np.integer] diff --git a/tests/series/test_series.py b/tests/series/test_series.py index 40c3ec10a..4e412133e 100644 --- a/tests/series/test_series.py +++ b/tests/series/test_series.py @@ -47,6 +47,7 @@ ) import xarray as xr +from pandas._libs.tslibs.offsets import Day from pandas._typing import ( DtypeObj, Scalar, @@ -61,7 +62,15 @@ check, ensure_clean, np_1darray, + np_1darray_anyint, np_1darray_bool, + np_1darray_bytes, + np_1darray_complex, + np_1darray_dt, + np_1darray_float, + np_1darray_object, + np_1darray_str, + np_1darray_td, np_ndarray_num, pytest_warns_bounded, ) @@ -1994,16 +2003,94 @@ def test_dtype_type() -> None: def test_types_to_numpy() -> None: s = pd.Series(["a", "b", "c"], dtype=str) - check(assert_type(s.to_numpy(), np_1darray), np_1darray) - check(assert_type(s.to_numpy(dtype="str", copy=True), np_1darray), np_1darray) - check(assert_type(s.to_numpy(na_value=0), np_1darray), np_1darray) - check(assert_type(s.to_numpy(na_value=np.int32(4)), np_1darray), np_1darray) - check(assert_type(s.to_numpy(na_value=np.float16(4)), np_1darray), np_1darray) - check(assert_type(s.to_numpy(na_value=np.complex128(4, 7)), np_1darray), np_1darray) + check(assert_type(s.to_numpy(), np_1darray[np.str_]), np_1darray) + check( + assert_type(s.to_numpy(dtype="str", copy=True), np_1darray[np.str_]), np_1darray + ) + check(assert_type(s.to_numpy(na_value=0), np_1darray[np.str_]), np_1darray) + check( + assert_type(s.to_numpy(na_value=np.int32(4)), np_1darray[np.str_]), np_1darray + ) + check( + assert_type(s.to_numpy(na_value=np.float16(4)), np_1darray[np.str_]), np_1darray + ) + check( + assert_type(s.to_numpy(na_value=np.complex128(4, 7)), np_1darray[np.str_]), + np_1darray, + ) check(assert_type(pd.Series().to_numpy(), np_1darray), np_1darray) +def test_to_numpy() -> None: + """Test Series.to_numpy for different types.""" + s_str = pd.Series(["a", "b", "c"], dtype=str) + check(assert_type(s_str.to_numpy(), np_1darray_str), np_1darray, str) + + s_bytes = pd.Series(["a", "b", "c"]).astype(bytes) + check(assert_type(s_bytes.to_numpy(), np_1darray_bytes), np_1darray, np.bytes_) + + s_bool = pd.Series([True, False]) + check(assert_type(s_bool.to_numpy(), np_1darray_bool), np_1darray, np.bool_) + + s_int = pd.Series([2, 3, 4]) + check(assert_type(s_int.to_numpy(), np_1darray_anyint), np_1darray, np.integer) + + s_float = pd.Series([2.0, 3.54, 4.84]) + check( + assert_type(s_float.to_numpy(), np_1darray_float), + np_1darray, + np.floating, + ) + + s_complex = pd.Series([2.0 + 2j, 3.54 + 4j, 4.84]) + check( + assert_type(s_complex.to_numpy(), np_1darray_complex), + np_1darray, + np.complexfloating, + ) + + dates = pd.Series( + [ + pd.Timestamp("2020-01-01"), + pd.Timestamp("2020-01-15"), + pd.Timestamp("2020-02-01"), + ], + dtype="datetime64[ns]", + ) + s_period = pd.PeriodIndex(dates, freq="M").to_series() + check(assert_type(s_period.to_numpy(), np_1darray_object), np_1darray, pd.Period) + + s_interval = pd.Series( + [ + pd.Interval(date, date + pd.DateOffset(days=1), closed="left") + for date in dates + ] + ) + check( + assert_type(s_interval.to_numpy(), np_1darray_object), np_1darray, pd.Interval + ) + + s_day = pd.Series([Day(1)]) + check(assert_type(s_day.to_numpy(), np_1darray_object), np_1darray, Day) + + s_date = pd.Series(pd.date_range(start="2017-01-01", end="2017-02-01")) + check( + assert_type(s_date.to_numpy(), np_1darray_dt), + np_1darray, + np.datetime64, + ) + + s_timedelta = pd.Series( + [pd.Timestamp.now().date(), pd.Timestamp.now().date()] + ).diff() + check( + assert_type(s_timedelta.to_numpy(), np_1darray_td), + np_1darray, + np.timedelta64, + ) + + def test_where() -> None: s = pd.Series([1, 2, 3], dtype=int)