Support DataFrame.astype('category') (pandas-dev#559)

* Add test for Series as CategoricalDtype

* Add `CategoryDtypeArg` to `AstypeArg`

* Apply review suggestions

* Add test for `DataFrame.astype('category')`

Author: zmoon

pandas-stubs/_typing.pyi

@@ -148,7 +148,7 @@ ComplexDtypeArg: TypeAlias = (
 )
 TimedeltaDtypeArg: TypeAlias = Literal["timedelta64[ns]"]
 TimestampDtypeArg: TypeAlias = Literal["datetime64[ns]"]
-CategoryDtypeArg: TypeAlias = Literal["category"]
+CategoryDtypeArg: TypeAlias = CategoricalDtype | Literal["category"]
 
 AstypeArg: TypeAlias = (
     BooleanDtypeArg
@@ -159,7 +159,7 @@ AstypeArg: TypeAlias = (
     | ComplexDtypeArg
     | TimedeltaDtypeArg
     | TimestampDtypeArg
-    | CategoricalDtype
+    | CategoryDtypeArg
     | ExtensionDtype
 )
 # DtypeArg specifies all allowable dtypes in a functions its dtype argument

tests/test_frame.py

@@ -2442,3 +2442,5 @@ def test_astype() -> None:
         assert_type(ab.astype({"col1": "int32", "col2": str}), "pd.DataFrame"),
         pd.DataFrame,
     )
+    check(assert_type(s.astype(pd.CategoricalDtype()), "pd.DataFrame"), pd.DataFrame)
+    check(assert_type(s.astype("category"), "pd.DataFrame"), pd.DataFrame)  # GH 559

tests/test_series.py

@@ -1591,3 +1591,15 @@ def test_updated_astype() -> None:
         pd.Series,
         Decimal,
     )
+
+    # Categorical
+    check(
+        assert_type(s.astype(pd.CategoricalDtype()), "pd.Series[Any]"),
+        pd.Series,
+        np.integer,
+    )
+    check(
+        assert_type(s.astype("category"), "pd.Series[Any]"),
+        pd.Series,
+        np.integer,
+    )