added test for new dtypes
@@ -106,3 +106,61 @@ 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]
+
+
+def test_sanitize_nullable_integers():
+ try:
+ 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"))],
+ }
+ )
+ except TypeError: # dtype not supported
+ return
+
+ 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],
+ }
+
+
+def test_sanitize_string_dtype():
+ try:
+ 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"),
+ }
+ )
+ except TypeError: # dtype not supported
+ return
+
+ 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"],
+ }
