BUG: astype(Int64) raises AttributeError

Author: sinhrks

pandas/core/generic.py

@@ -28,6 +28,7 @@
     is_re_compilable,
     is_period_arraylike,
     is_object_dtype,
+    is_extension_array_dtype,
     pandas_dtype)
 from pandas.core.dtypes.cast import maybe_promote, maybe_upcast_putmask
 from pandas.core.dtypes.inference import is_hashable
@@ -5293,8 +5294,9 @@ def astype(self, dtype, copy=True, errors='raise', **kwargs):
                 else:
                     results.append(results.append(col.copy() if copy else col))
 
-        elif is_categorical_dtype(dtype) and self.ndim > 1:
+        elif is_extension_array_dtype(dtype) and self.ndim > 1:
             # GH 18099: columnwise conversion to categorical
+            # and extension dtype
             results = (self[col].astype(dtype, copy=copy) for col in self)
 
         else:

pandas/core/internals/blocks.py

@@ -671,15 +671,6 @@ def _astype(self, dtype, copy=False, errors='raise', values=None,
                 raise
             newb = self.copy() if copy else self
 
-        if newb.is_numeric and self.is_numeric:
-            if newb.shape != self.shape:
-                raise TypeError(
-                    "cannot set astype for copy = [{copy}] for dtype "
-                    "({dtype} [{itemsize}]) with smaller itemsize than "
-                    "current ({newb_dtype} [{newb_size}])".format(
-                        copy=copy, dtype=self.dtype.name,
-                        itemsize=self.itemsize, newb_dtype=newb.dtype.name,
-                        newb_size=newb.itemsize))
         return newb
 
     def convert(self, copy=True, **kwargs):

pandas/tests/frame/test_dtypes.py

@@ -12,6 +12,7 @@
 from pandas.compat import u
 from pandas import _np_version_under1p14
 
+from pandas.core.arrays import integer_array
 from pandas.core.dtypes.dtypes import DatetimeTZDtype, CategoricalDtype
 from pandas.tests.frame.common import TestData
 from pandas.util.testing import (assert_series_equal,
@@ -666,6 +667,48 @@ def test_astype_categoricaldtype_class_raises(self, cls):
         with tm.assert_raises_regex(TypeError, xpr):
             df['A'].astype(cls)
 
+    @pytest.mark.parametrize("dtype", ['Int64', 'Int32', 'Int16'])
+    def test_astype_extension_dtypes(self, dtype):
+        # GH 22578
+        df = pd.DataFrame([[1., 2.], [3., 4.], [5., 6.]], columns=['a', 'b'])
+
+        expected1 = pd.DataFrame({'a': integer_array([1, 3, 5],
+                                                     dtype=dtype),
+                                  'b': integer_array([2, 4, 6],
+                                                     dtype=dtype)})
+        tm.assert_frame_equal(df.astype(dtype), expected1)
+        tm.assert_frame_equal(df.astype('int64').astype(dtype), expected1)
+        tm.assert_frame_equal(df.astype(dtype).astype('float64'), df)
+
+        df = pd.DataFrame([[1., 2.], [3., 4.], [5., 6.]], columns=['a', 'b'])
+        df['b'] = df['b'].astype(dtype)
+        expected2 = pd.DataFrame({'a': [1., 3., 5.],
+                                  'b': integer_array([2, 4, 6],
+                                                     dtype=dtype)})
+        tm.assert_frame_equal(df, expected2)
+
+        tm.assert_frame_equal(df.astype(dtype), expected1)
+        tm.assert_frame_equal(df.astype('int64').astype(dtype), expected1)
+
+    @pytest.mark.parametrize("dtype", ['Int64', 'Int32', 'Int16'])
+    def test_astype_extension_dtypes_1d(self, dtype):
+        # GH 22578
+        df = pd.DataFrame({'a': [1., 2., 3.]})
+
+        expected1 = pd.DataFrame({'a': integer_array([1, 2, 3],
+                                                     dtype=dtype)})
+        tm.assert_frame_equal(df.astype(dtype), expected1)
+        tm.assert_frame_equal(df.astype('int64').astype(dtype), expected1)
+
+        df = pd.DataFrame({'a': [1., 2., 3.]})
+        df['a'] = df['a'].astype(dtype)
+        expected2 = pd.DataFrame({'a': integer_array([1, 2, 3],
+                                                     dtype=dtype)})
+        tm.assert_frame_equal(df, expected2)
+
+        tm.assert_frame_equal(df.astype(dtype), expected1)
+        tm.assert_frame_equal(df.astype('int64').astype(dtype), expected1)
+
     @pytest.mark.parametrize('dtype', [
         {100: 'float64', 200: 'uint64'}, 'category', 'float64'])
     def test_astype_column_metadata(self, dtype):