From da3852845abdc389ab12b1f4ac00bce03877ff83 Mon Sep 17 00:00:00 2001 From: Brock Date: Fri, 7 Jan 2022 18:53:57 -0800 Subject: [PATCH 1/2] BUG: DataFrame.astype preserve subclasses GH#40810 --- doc/source/whatsnew/v1.5.0.rst | 2 +- pandas/core/generic.py | 1 + pandas/tests/frame/test_subclass.py | 7 +++++++ 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/doc/source/whatsnew/v1.5.0.rst b/doc/source/whatsnew/v1.5.0.rst index 2cbc7b06b89df..2636bf61c993a 100644 --- a/doc/source/whatsnew/v1.5.0.rst +++ b/doc/source/whatsnew/v1.5.0.rst @@ -138,7 +138,7 @@ Numeric Conversion ^^^^^^^^^^ -- +- Bug in :meth:`DataFrame.astype` not preserving subclasses (:issue:`40810`) - Strings diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 193be98e904b9..779a8666f8616 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -5891,6 +5891,7 @@ def astype( # GH 19920: retain column metadata after concat result = concat(results, axis=1, copy=False) + result = self._constructor(result) # GH#40810 retain subclass result.columns = self.columns result = result.__finalize__(self, method="astype") # https://github.com/python/mypy/issues/8354 diff --git a/pandas/tests/frame/test_subclass.py b/pandas/tests/frame/test_subclass.py index 8d9957b24300f..0d6cf39f801db 100644 --- a/pandas/tests/frame/test_subclass.py +++ b/pandas/tests/frame/test_subclass.py @@ -723,6 +723,13 @@ def test_convert_dtypes_preserves_subclass(self, gpd_style_subclass_df): result = gpd_style_subclass_df.convert_dtypes() assert isinstance(result, type(gpd_style_subclass_df)) + def test_astype_preserves_subclass(self): + # GH#40810 + df = tm.SubclassedDataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]}) + + result = df.astype({"A": np.int64, "B": np.int32, "C": np.float64}) + assert isinstance(result, tm.SubclassedDataFrame) + def test_equals_subclass(self): # https://github.com/pandas-dev/pandas/pull/34402 # allow subclass in both directions From 4865822efe41fee27737cdcb291bca91b289aeaa Mon Sep 17 00:00:00 2001 From: Brock Date: Sat, 8 Jan 2022 14:52:22 -0800 Subject: [PATCH 2/2] mypy fixup --- pandas/core/generic.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 779a8666f8616..891e1b77534d0 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -5891,7 +5891,12 @@ def astype( # GH 19920: retain column metadata after concat result = concat(results, axis=1, copy=False) - result = self._constructor(result) # GH#40810 retain subclass + # GH#40810 retain subclass + # Incompatible types in assignment (expression has type "NDFrameT", + # variable has type "DataFrame") + # Argument 1 to "NDFrame" has incompatible type "DataFrame"; expected + # "Union[ArrayManager, SingleArrayManager, BlockManager, SingleBlockManager]" + result = self._constructor(result) # type: ignore[arg-type,assignment] result.columns = self.columns result = result.__finalize__(self, method="astype") # https://github.com/python/mypy/issues/8354