Add type parameter to remaining Series.astype overloads. (#949)

* Add type parameter to remaining Series.astype overloads. * Make on overload more precise. * Fix CategoricalDtype
pandas-dev · Jun 28, 2024 · e1ba6ad · e1ba6ad
1 parent ea8bdee
commit e1ba6ad
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
@@ -158,6 +158,7 @@ from pandas._typing import (
 )
 
 from pandas.core.dtypes.base import ExtensionDtype
+from pandas.core.dtypes.dtypes import CategoricalDtype
 
 from pandas.plotting import PlotAccessor
 
@@ -1320,14 +1321,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
         dtype: CategoryDtypeArg,
         copy: _bool = ...,
         errors: IgnoreRaise = ...,
-    ) -> Series: ...
+    ) -> Series[CategoricalDtype]: ...
     @overload
     def astype(
         self,
         dtype: ObjectDtypeArg | VoidDtypeArg | ExtensionDtype | DtypeObj,
         copy: _bool = ...,
         errors: IgnoreRaise = ...,
-    ) -> Series: ...
+    ) -> Series[Any]: ...
     def copy(self, deep: _bool = ...) -> Series[S1]: ...
     def infer_objects(self) -> Series[S1]: ...
     @overload

diff --git a/tests/test_series.py b/tests/test_series.py
@@ -2581,8 +2581,8 @@ def test_astype_categorical(cast_arg: CategoryDtypeArg, target_type: type) -> No
 
     if TYPE_CHECKING:
         # pandas category
-        assert_type(s.astype(pd.CategoricalDtype()), "pd.Series[Any]")
-        assert_type(s.astype("category"), "pd.Series[Any]")
+        assert_type(s.astype(pd.CategoricalDtype()), "pd.Series[pd.CategoricalDtype]")
+        assert_type(s.astype("category"), "pd.Series[pd.CategoricalDtype]")
         # pyarrow dictionary
         # assert_type(s.astype("dictionary[pyarrow]"), "pd.Series[Categorical]")