gh-565: Added Object to Astype (#566)

* added object support to astype * Update test.py * Delete test.py
pandas-dev · Mar 9, 2023 · 927a2d0 · 927a2d0
1 parent 6829a99
commit 927a2d0
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 1 deletion.
diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi
@@ -192,6 +192,7 @@ AstypeArg: TypeAlias = (
     | TimestampDtypeArg
     | CategoryDtypeArg
     | ExtensionDtype
+    | type[object]
 )
 # DtypeArg specifies all allowable dtypes in a functions its dtype argument
 DtypeObj: TypeAlias = np.dtype[np.generic] | ExtensionDtype

diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
@@ -1118,7 +1118,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     @overload
     def astype(
         self,
-        dtype: ExtensionDtype,
+        dtype: type[object] | ExtensionDtype,
         copy: _bool = ...,
         errors: IgnoreRaise = ...,
     ) -> Series: ...

diff --git a/tests/test_frame.py b/tests/test_frame.py
@@ -2447,3 +2447,23 @@ def test_astype() -> None:
     )
     check(assert_type(s.astype(pd.CategoricalDtype()), "pd.DataFrame"), pd.DataFrame)
     check(assert_type(s.astype("category"), "pd.DataFrame"), pd.DataFrame)  # GH 559
+
+    population_dict = {
+        "California": 38332521,
+        "Texas": 26448193,
+        "New York": 19651127,
+        "Florida": 19552860,
+        "Illinois": 12882135,
+    }
+    area_dict = {
+        "California": 423967,
+        "Texas": 695662,
+        "New York": 141297,
+        "Florida": 170312,
+        "Illinois": 149995,
+    }
+    population = pd.Series(population_dict)
+    area = pd.Series(area_dict)
+
+    states = pd.DataFrame({"population": population, "area": area})
+    check(assert_type(states.astype(object), pd.DataFrame), pd.DataFrame, object)
diff --git a/tests/test_series.py b/tests/test_series.py
@@ -1722,6 +1722,21 @@ def test_updated_astype() -> None:
         Decimal,
     )
 
+    s4 = pd.Series([1, 1])
+    s5 = pd.Series([s4, 4])
+    population_dict = {
+        "California": 38332521,
+        "Texas": 26448193,
+        "New York": 19651127,
+        "Florida": 19552860,
+        "Illinois": 12882135,
+    }
+    population = pd.Series(population_dict)
+
+    check(assert_type(s4.astype(object), pd.Series), pd.Series, object)
+    check(assert_type(s5.astype(object), pd.Series), pd.Series, object)
+    check(assert_type(population.astype(object), pd.Series), pd.Series, object)
+
     # Categorical
     check(
         assert_type(s.astype(pd.CategoricalDtype()), "pd.Series[Any]"),