Repository: altair-viz/altair
Pull request #1924 (2020-01-21 10:48:08)

Nullable integers and string dtype
Add support for the new dtypes in `utils.sanitize_dataframe` to solve #1922

- [nullable integer datatype](https://pandas.pydata.org/pandas-docs/version/0.25/whatsnew/v0.24.0.html#optional-integer-na-support) 
- [dedicated string dtype](https://pandas.pydata.org/pandas-docs/version/1.0.0/whatsnew/v1.0.0.html#dedicated-string-data-type)




Patch:
diff --git a/altair/utils/core.py b/altair/utils/core.py
index 73b1cfea..25cb72c6 100644
--- a/altair/utils/core.py
+++ b/altair/utils/core.py
@@ -171,6 +171,9 @@ def sanitize_dataframe(df):
     * Convert np.int dtypes to Python int objects
     * Convert floats to objects and replace NaNs/infs with None.
     * Convert DateTime dtypes into appropriate string representations
+    * Convert Nullable integers to objects and replace NaN with None
+    * Convert Nullable boolean to objects and replace NaN with None
+    * convert dedicated string column to objects and replace NaN with None
     * Raise a ValueError for TimeDelta dtypes
     """
     df = df.copy()
@@ -200,9 +203,19 @@ def sanitize_dataframe(df):
             # https://github.com/pydata/pandas/issues/10778
             col = df[col_name].astype(object)
             df[col_name] = col.where(col.notnull(), None)
+        elif str(dtype) == "string":
+            # dedicated string datatype (since 1.0)
+            # https://pandas.pydata.org/pandas-docs/version/1.0.0/whatsnew/v1.0.0.html#dedicated-string-data-type
+            col = df[col_name].astype(object)
+            df[col_name] = col.where(col.notnull(), None)
         elif str(dtype) == 'bool':
             # convert numpy bools to objects; np.bool is not JSON serializable
             df[col_name] = df[col_name].astype(object)
+        elif str(dtype) == "boolean":
+            # dedicated boolean datatype (since 1.0)
+            # https://pandas.io/docs/user_guide/boolean.html
+            col = df[col_name].astype(object)
+            df[col_name] = col.where(col.notnull(), None)
         elif str(dtype).startswith('datetime'):
             # Convert datetimes to strings. This needs to be a full ISO string
             # with time, which is why we cannot use ``col.astype(str)``.
@@ -220,6 +233,19 @@ def sanitize_dataframe(df):
             # geopandas >=0.6.1 uses the dtype geometry. Continue here
             # otherwise it will give an error on np.issubdtype(dtype, np.integer)
             continue                            
+        elif str(dtype) in {
+            "Int8",
+            "Int16",
+            "Int32",
+            "Int64",
+            "UInt8",
+            "UInt16",
+            "UInt32",
+            "UInt64",
+        }:  # nullable integer datatypes (since 24.0)
+            # https://pandas.pydata.org/pandas-docs/version/0.25/whatsnew/v0.24.0.html#optional-integer-na-support
+            col = df[col_name].astype(object)
+            df[col_name] = col.where(col.notnull(), None)
         elif np.issubdtype(dtype, np.integer):
             # convert integers to objects; np.int is not JSON serializable
             df[col_name] = df[col_name].astype(object)


Test patch:
diff --git a/altair/utils/tests/test_utils.py b/altair/utils/tests/test_utils.py
index d8851913..3f92afed 100644
--- a/altair/utils/tests/test_utils.py
+++ b/altair/utils/tests/test_utils.py
@@ -106,3 +106,90 @@ def test_sanitize_dataframe_infs():
     df_clean = sanitize_dataframe(df)
     assert list(df_clean.dtypes) == [object]
     assert list(df_clean['x']) == [0, 1, 2, None, None, None]
+
+
+@pytest.mark.skipif(
+    not hasattr(pd, "Int64Dtype"),
+    reason="Nullable integers not supported in pandas v{}".format(pd.__version__)
+    )
+def test_sanitize_nullable_integers():
+
+    df = pd.DataFrame(
+        {
+            "int_np": [1, 2, 3, 4, 5],
+            "int64": pd.Series([1, 2, 3, None, 5], dtype="UInt8"),
+            "int64_nan": pd.Series([1, 2, 3, float("nan"), 5], dtype="Int64"),
+            "float": [1.0, 2.0, 3.0, 4, 5.0],
+            "float_null": [1, 2, None, 4, 5],
+            "float_inf": [1, 2, None, 4, (float("inf"))],
+        }
+    )
+
+    df_clean = sanitize_dataframe(df)
+    assert {col.dtype.name for _, col in df_clean.iteritems()} == {"object"}
+
+    result_python = {
+        col_name: list(col) for col_name, col in df_clean.iteritems()
+    }
+    assert result_python == {
+        "int_np": [1, 2, 3, 4, 5],
+        "int64": [1, 2, 3, None, 5],
+        "int64_nan": [1, 2, 3, None, 5],
+        "float": [1.0, 2.0, 3.0, 4.0, 5.0],
+        "float_null": [1.0, 2.0, None, 4.0, 5.0],
+        "float_inf": [1.0, 2.0, None, 4.0, None],
+    }
+
+
+@pytest.mark.skipif(
+    not hasattr(pd, "StringDtype"),
+    reason="dedicated String dtype not supported in pandas v{}".format(pd.__version__)
+    )
+def test_sanitize_string_dtype():
+    df = pd.DataFrame(
+        {
+            "string_object": ["a", "b", "c", "d"],
+            "string_string": pd.array(["a", "b", "c", "d"], dtype="string"),
+            "string_object_null": ["a", "b", None, "d"],
+            "string_string_null": pd.array(["a", "b", None, "d"], dtype="string"),
+        }
+    )
+
+    df_clean = sanitize_dataframe(df)
+    assert {col.dtype.name for _, col in df_clean.iteritems()} == {"object"}
+
+    result_python = {
+        col_name: list(col) for col_name, col in df_clean.iteritems()
+    }
+    assert result_python == {
+        "string_object": ["a", "b", "c", "d"],
+        "string_string": ["a", "b", "c", "d"],
+        "string_object_null": ["a", "b", None, "d"],
+        "string_string_null": ["a", "b", None, "d"],
+    }
+
+
+@pytest.mark.skipif(
+    not hasattr(pd, "BooleanDtype"),
+    reason="Nullable boolean dtype not supported in pandas v{}".format(pd.__version__)
+    )
+def test_sanitize_boolean_dtype():
+    df = pd.DataFrame(
+        {
+            "bool_none": pd.array([True, False, None], dtype="boolean"),
+            "none": pd.array([None, None, None], dtype="boolean"),
+            "bool": pd.array([True, False, True], dtype="boolean"),
+        }
+    )
+
+    df_clean = sanitize_dataframe(df)
+    assert {col.dtype.name for _, col in df_clean.iteritems()} == {"object"}
+
+    result_python = {
+        col_name: list(col) for col_name, col in df_clean.iteritems()
+    }
+    assert result_python == {
+            "bool_none": [True, False, None],
+            "none": [None, None, None],
+            "bool": [True, False, True],
+    }