TST(string dtype): Resolve HDF5 xfails in test_put.py

Author: rhshadrach

pandas/io/pytables.py

@@ -86,6 +86,7 @@
     PeriodArray,
 )
 from pandas.core.arrays.datetimes import tz_to_dtype
+from pandas.core.arrays.string_ import BaseStringArray
 import pandas.core.common as com
 from pandas.core.computation.pytables import (
     PyTablesExpr,
@@ -3185,6 +3186,8 @@ def write_array(
         #  both self._filters and EA
 
         value = extract_array(obj, extract_numpy=True)
+        if isinstance(value, BaseStringArray):
+            value = value.to_numpy()
 
         if key in self.group:
             self._handle.remove_node(self.group, key)
@@ -3363,7 +3366,11 @@ def read(
 
             columns = items[items.get_indexer(blk_items)]
             df = DataFrame(values.T, columns=columns, index=axes[1], copy=False)
-            if using_string_dtype() and is_string_array(values, skipna=True):
+            if (
+                using_string_dtype()
+                and isinstance(values, np.ndarray)
+                and is_string_array(values, skipna=True)
+            ):
                 df = df.astype(StringDtype(na_value=np.nan))
             dfs.append(df)
 
@@ -4737,9 +4744,10 @@ def read(
                 df = DataFrame._from_arrays([values], columns=cols_, index=index_)
             if not (using_string_dtype() and values.dtype.kind == "O"):
                 assert (df.dtypes == values.dtype).all(), (df.dtypes, values.dtype)
-            if using_string_dtype() and is_string_array(
-                values,  # type: ignore[arg-type]
-                skipna=True,
+            if (
+                using_string_dtype()
+                and isinstance(values, np.ndarray)
+                and is_string_array(values, skipna=True)
             ):
                 df = df.astype(StringDtype(na_value=np.nan))
             frames.append(df)

pandas/tests/io/pytables/test_put.py

@@ -3,8 +3,6 @@
 import numpy as np
 import pytest
 
-from pandas._config import using_string_dtype
-
 from pandas._libs.tslibs import Timestamp
 
 import pandas as pd
@@ -26,7 +24,6 @@
 
 pytestmark = [
     pytest.mark.single_cpu,
-    pytest.mark.xfail(using_string_dtype(), reason="TODO(infer_string)", strict=False),
 ]
 
 
@@ -99,7 +96,7 @@ def test_api_default_format(tmp_path, setup_path):
             assert store.get_storer("df4").is_table
 
 
-def test_put(setup_path):
+def test_put(setup_path, using_infer_string):
     with ensure_clean_store(setup_path) as store:
         ts = Series(
             np.arange(10, dtype=np.float64), index=date_range("2020-01-01", periods=10)
@@ -133,7 +130,11 @@ def test_put(setup_path):
 
         # overwrite table
         store.put("c", df[:10], format="table", append=False)
-        tm.assert_frame_equal(df[:10], store["c"])
+        expected = df[:10]
+        if using_infer_string:
+            expected.columns = expected.columns.astype("str")
+        result = store["c"]
+        tm.assert_frame_equal(result, expected)
 
 
 def test_put_string_index(setup_path):
@@ -162,7 +163,7 @@ def test_put_string_index(setup_path):
         tm.assert_frame_equal(store["b"], df)
 
 
-def test_put_compression(setup_path):
+def test_put_compression(setup_path, using_infer_string):
     with ensure_clean_store(setup_path) as store:
         df = DataFrame(
             np.random.default_rng(2).standard_normal((10, 4)),
@@ -171,7 +172,11 @@ def test_put_compression(setup_path):
         )
 
         store.put("c", df, format="table", complib="zlib")
-        tm.assert_frame_equal(store["c"], df)
+        expected = df
+        if using_infer_string:
+            expected.columns = expected.columns.astype("str")
+        result = store["c"]
+        tm.assert_frame_equal(result, expected)
 
         # can't compress if format='fixed'
         msg = "Compression not supported on Fixed format stores"
@@ -180,7 +185,7 @@ def test_put_compression(setup_path):
 
 
 @td.skip_if_windows
-def test_put_compression_blosc(setup_path):
+def test_put_compression_blosc(setup_path, using_infer_string):
     df = DataFrame(
         np.random.default_rng(2).standard_normal((10, 4)),
         columns=Index(list("ABCD"), dtype=object),
@@ -194,10 +199,14 @@ def test_put_compression_blosc(setup_path):
             store.put("b", df, format="fixed", complib="blosc")
 
         store.put("c", df, format="table", complib="blosc")
-        tm.assert_frame_equal(store["c"], df)
+        expected = df
+        if using_infer_string:
+            expected.columns = expected.columns.astype("str")
+        result = store["c"]
+        tm.assert_frame_equal(result, expected)
 
 
-def test_put_mixed_type(setup_path, performance_warning):
+def test_put_mixed_type(setup_path, performance_warning, using_infer_string):
     df = DataFrame(
         np.random.default_rng(2).standard_normal((10, 4)),
         columns=Index(list("ABCD"), dtype=object),
@@ -223,8 +232,11 @@ def test_put_mixed_type(setup_path, performance_warning):
         with tm.assert_produces_warning(performance_warning):
             store.put("df", df)
 
-        expected = store.get("df")
-        tm.assert_frame_equal(expected, df)
+        expected = df
+        if using_infer_string:
+            expected.columns = expected.columns.astype("str")
+        result = store.get("df")
+        tm.assert_frame_equal(result, expected)
 
 
 @pytest.mark.parametrize("format", ["table", "fixed"])
@@ -253,7 +265,7 @@ def test_store_index_types(setup_path, format, index):
         tm.assert_frame_equal(df, store["df"])
 
 
-def test_column_multiindex(setup_path):
+def test_column_multiindex(setup_path, using_infer_string):
     # GH 4710
     # recreate multi-indexes properly
 
@@ -264,6 +276,11 @@ def test_column_multiindex(setup_path):
     expected = df.set_axis(df.index.to_numpy())
 
     with ensure_clean_store(setup_path) as store:
+        if using_infer_string:
+            msg = "Saving a MultiIndex with an extension dtype is not supported."
+            with pytest.raises(NotImplementedError, match=msg):
+                store.put("df", df)
+            return
         store.put("df", df)
         tm.assert_frame_equal(
             store["df"], expected, check_index_type=True, check_column_type=True