ENH: Implement convert_dtypes

doc/source/user_guide/missing_data.rst

@@ -806,7 +806,8 @@ dtype, it will use ``pd.NA``:
 
 Currently, pandas does not yet use those data types by default (when creating
 a DataFrame or Series, or when reading in data), so you need to specify
-the dtype explicitly.
+the dtype explicitly.  An easy way to convert to those dtypes is explained
+:ref:`here <missing_data.NA.conversion>`.
 
 Propagation in arithmetic and comparison operations
 ---------------------------------------------------
@@ -945,3 +946,25 @@ work with ``NA``, and generally return ``NA``:
    in the future.
 
 See :ref:`dsintro.numpy_interop` for more on ufuncs.
+
+.. _missing_data.NA.conversion:
+
+Conversion
+----------
+
+If you have a DataFrame or Series using traditional types that have missing data
+represented using ``np.nan``, there are convenience methods
+:meth:`~Series.as_nullable_dtypes` in Series and :meth:`~DataFrame.as_nullable_dtypes`
+in DataFrame that can convert data to use the newer dtypes for integers, strings and
+booleans listed :ref:`here <basics.dtypes>`. This is especially helpful after reading
+in data sets when letting the readers infer default dtypes.
+
+In this example, while the dtypes of all columns are changed, we show the results for
+the first 10 columns.
+
+.. ipython:: python
+
+   bb = pd.read_csv('data/baseball.csv', index_col='id')
+   bb[bb.columns[:10]].dtypes
+   bbn = bb.as_nullable_dtypes()
+   bbn[bbn.columns[:10]].dtypes

doc/source/whatsnew/v1.0.0.rst

@@ -228,6 +228,7 @@ Other enhancements
 - Added an experimental :attr:`~DataFrame.attrs` for storing global metadata about a dataset (:issue:`29062`)
 - :meth:`Timestamp.fromisocalendar` is now compatible with python 3.8 and above (:issue:`28115`)
 - :meth:`DataFrame.to_pickle` and :func:`read_pickle` now accept URL (:issue:`30163`)
+- Added :meth:`DataFrame.as_nullable_dtypes` and :meth:`Series.as_nullable_dtypes` to make it easier to use ``pd.NA``. See :ref:`here <missing_data.NA.conversion>` for a description. (:issue:`29752`)
 
 
 Build Changes

pandas/core/generic.py

@@ -5909,6 +5909,97 @@ def infer_objects(self: FrameOrSeries) -> FrameOrSeries:
             )
         ).__finalize__(self)
 
+    def as_nullable_dtypes(
+        self: FrameOrSeries, convert_integer: bool_t = True
+    ) -> FrameOrSeries:
+        """
+        Convert columns of DataFrame or a Series to types supporting ``pd.NA``.
+
+        If the dtype is ``object``, if possible, convert to ``StringDtype``,
+        ``BooleanDtype`` or an appropriate integer extension type, otherwise leave as
+        ``object``.
+
+        If the dtype is integer", convert to an appropriate integer type.
+
+        If the dtype is numeric, and consists of all integers, convert to an
+        appropriate integer extension type.
+
+        Parameters
+        ----------
+        convert_integer : bool, default True
+            Whether ``int`` types should be converted to integer extension types
+
+        Returns
+        -------
+        converted : same type as input object
+
+        Examples
+        --------
+        >>> df = pd.DataFrame(
+        ...     {
+        ...         "a": pd.Series([1, 2, 3], dtype=np.dtype("int32")),
+        ...         "b": pd.Series(["x", "y", "z"], dtype=np.dtype("O")),
+        ...         "c": pd.Series([True, False, np.nan], dtype=np.dtype("O")),
+        ...         "d": pd.Series(["h", "i", np.nan], dtype=np.dtype("O")),
+        ...         "e": pd.Series([10, np.nan, 20], dtype=np.dtype("float")),
+        ...         "f": pd.Series([np.nan, 100.5, 200], dtype=np.dtype("float")),
+        ...     }
+        ... )
+
+        >>> df
+           a  b      c    d     e      f
+        0  1  x   True    h  10.0    NaN
+        1  2  y  False    i   NaN  100.5
+        2  3  z    NaN  NaN  20.0  200.0
+
+        >>> df.dtypes
+        a      int32
+        b     object
+        c     object
+        d     object
+        e    float64
+        f    float64
+        dtype: object
+
+        >>> dfn = df.as_nullable_dtypes()
+        >>> dfn
+           a  b      c     d     e      f
+        0  1  x   True     h    10    NaN
+        1  2  y  False     i  <NA>  100.5
+        2  3  z   <NA>  <NA>    20  200.0
+
+        >>> dfn.dtypes
+        a      Int32
+        b     string
+        c    boolean
+        d     string
+        e      Int64
+        f    float64
+        dtype: object
+
+        >>> s = pd.Series(["a", "b", np.nan])
+        >>> s
+        0      a
+        1      b
+        2    NaN
+        dtype: object
+
+        >>> s.as_nullable_dtypes()
+        0       a
+        1       b
+        2    <NA>
+        dtype: string
+        """
+        if self.ndim == 1:
+            return self._as_nullable_dtype(convert_integer)
+        else:
+            results = [
+                col._as_nullable_dtype(convert_integer)
+                for col_name, col in self.items()
+            ]
+            result = pd.concat(results, axis=1, copy=False)
+            return result
+
     # ----------------------------------------------------------------------
     # Filling NA's
 

pandas/core/series.py

@@ -37,8 +37,10 @@
     is_dict_like,
     is_extension_array_dtype,
     is_integer,
+    is_integer_dtype,
     is_iterator,
     is_list_like,
+    is_numeric_dtype,
     is_object_dtype,
     is_scalar,
     is_timedelta64_dtype,
@@ -62,6 +64,7 @@
 from pandas.core.accessor import CachedAccessor
 from pandas.core.arrays import ExtensionArray, try_cast_to_ea
 from pandas.core.arrays.categorical import Categorical, CategoricalAccessor
+from pandas.core.arrays.integer import _dtypes
 from pandas.core.arrays.sparse import SparseAccessor
 import pandas.core.common as com
 from pandas.core.construction import (
@@ -4332,6 +4335,65 @@ def between(self, left, right, inclusive=True) -> "Series":
 
         return lmask & rmask
 
+    # ----------------------------------------------------------------------
+    # Convert to types that support pd.NA
+
+    def _as_nullable_dtype(self: ABCSeries, convert_integer: bool = True) -> ABCSeries:
+        """
+        Convert columns of DataFrame or a Series to types supporting ``pd.NA``.
+
+        If the dtype is ``object``, if possible, convert to ``StringDtype``,
+        ``BooleanDtype`` or an appropriate integer extension type, otherwise leave as
+        ``object``.
+
+        If the dtype is "integer", convert to an appropriate integer type.
+
+        If the dtype is numeric, and consists of all integers, convert to an
+        appropriate integer type.
+
+        Parameters
+        ----------
+        convert_integer : bool, default True
+            Whether ``int`` types should be converted to integer extension types
+
+        Returns
+        -------
+        converted : same type as input object
+        """
+        target_int_dtype = "Int64"
+
+        try:
+            inferred_dtype = lib.infer_dtype(self)
+        except ValueError:
+            inferred_dtype = self.dtype
+
+        # If an object, try to convert to an integer, string or boolean
+        # extension type, otherwise leave it alone
+        if is_object_dtype(self.dtype):
+            if (
+                inferred_dtype == "mixed-integer"
+                or inferred_dtype == "mixed-integer-float"
+            ):
+                inferred_dtype = target_int_dtype
+            elif inferred_dtype not in {"string", "boolean", "integer"}:
+                inferred_dtype = self.dtype
+
+        # If an integer, then match the size based on the registry
+        elif is_integer_dtype(self.dtype):
+            if convert_integer and not is_extension_array_dtype(self.dtype):
+                inferred_dtype = _dtypes.get(self.dtype.name, target_int_dtype)
+
+        # If it's not integer and numeric try to make it an integer
+        elif is_numeric_dtype(self.dtype):
+            inferred_dtype = target_int_dtype
+
+        try:
+            result = self.astype(inferred_dtype)
+        except TypeError:
+            result = self.astype(self.dtype)
+
+        return result
+
     @Appender(generic._shared_docs["isna"] % _shared_doc_kwargs)
     def isna(self) -> "Series":
         return super().isna()

pandas/tests/frame/test_dtypes.py

@@ -1072,6 +1072,27 @@ def test_str_to_small_float_conversion_type(self):
         expected = pd.DataFrame(col_data, columns=["A"], dtype=float)
         tm.assert_frame_equal(result, expected)
 
+    @pytest.mark.parametrize(
+        "convert_integer, expected", [(False, np.dtype("int32")), (True, "Int32")]
+    )
+    def test_as_nullable_dtypes(self, convert_integer, expected):
+        # Specific types are tested in tests/series/test_dtypes.py
+        # Just check that it works for DataFrame here
+        df = pd.DataFrame(
+            {
+                "a": pd.Series([1, 2, 3], dtype=np.dtype("int32")),
+                "b": pd.Series(["x", "y", "z"], dtype=np.dtype("O")),
+            }
+        )
+        result = df.as_nullable_dtypes(convert_integer)
+        expected = pd.DataFrame(
+            {
+                "a": pd.Series([1, 2, 3], dtype=expected),
+                "b": pd.Series(["x", "y", "z"], dtype="string"),
+            }
+        )
+        tm.assert_frame_equal(result, expected)
+
 
 class TestDataFrameDatetimeWithTZ:
     def test_interleave(self, timezone_frame):

pandas/tests/series/test_dtypes.py

@@ -487,3 +487,47 @@ def test_reindex_astype_order_consistency(self):
         s1 = s.reindex(new_index).astype(temp_dtype).astype(new_dtype)
         s2 = s.astype(temp_dtype).reindex(new_index).astype(new_dtype)
         tm.assert_series_equal(s1, s2)
+
+    @pytest.mark.parametrize(
+        "series, dtype",
+        [
+            (Series([1, 2, 3], dtype=np.dtype("int32")), "Int32"),
+            (Series([1, 2, 3], dtype=np.dtype("int64")), "Int64"),
+            (Series(["x", "y", "z"], dtype=np.dtype("O")), pd.StringDtype()),
+            (Series([True, False, np.nan], dtype=np.dtype("O")), pd.BooleanDtype()),
+            (Series(["h", "i", np.nan], dtype=np.dtype("O")), pd.StringDtype()),
+            (Series([10, np.nan, 20], dtype=np.dtype("float")), pd.Int64Dtype()),
+            (Series([np.nan, 100.5, 200], dtype=np.dtype("float")), np.dtype("float")),
+            (Series([3, 4, 5], dtype="Int8"), "Int8"),
+            (Series([[1, 2], [3, 4], [5]]), np.dtype("O")),
+            (Series([4, 5, 6], dtype=np.dtype("uint32")), "UInt32"),
+            (Series([-10, 12, 13], dtype=np.dtype("i1")), "Int8"),
+            (Series([1, 2.0], dtype=object), "Int64"),
+            (Series(["a", "b"], dtype=pd.CategoricalDtype()), pd.CategoricalDtype()),
+            (
+                Series(
+                    pd.to_datetime(["2020-01-14 10:00", "2020-01-15 11:11"]),
+                    dtype=pd.DatetimeTZDtype(tz="UTC"),
+                ),
+                pd.DatetimeTZDtype(tz="UTC"),
+            ),
+            (
+                Series(
+                    pd.to_datetime(["2020-01-14 10:00", "2020-01-15 11:11"]),
+                    dtype="datetime64[ns]",
+                ),
+                np.dtype("datetime64[ns]"),
+            ),
+            (
+                Series(pd.period_range("1/1/2011", freq="M", periods=3)),
+                pd.PeriodDtype("M"),
+            ),
+            (
+                Series(pd.arrays.IntervalArray([pd.Interval(0, 1), pd.Interval(1, 5)])),
+                pd.IntervalDtype("int64"),
+            ),
+        ],
+    )
+    def test_as_nullable_dtypes(self, series, dtype):
+        ns = series.as_nullable_dtypes()
+        assert ns.dtype == dtype