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

Author: rhshadrach

pandas/io/pytables.py

@@ -1950,6 +1950,7 @@ def _write_to_group(
     def _read_group(self, group: Node):
         s = self._create_storer(group)
         s.infer_axes()
+        print(type(s), s)
         return s.read()
 
     def _identify_group(self, key: str, append: bool) -> Node:
@@ -3297,7 +3298,12 @@ def read(
         index = self.read_index("index", start=start, stop=stop)
         values = self.read_array("values", start=start, stop=stop)
         result = Series(values, index=index, name=self.name, copy=False)
-        if using_string_dtype() and is_string_array(values, skipna=True):
+        if (
+            using_string_dtype()
+            and isinstance(values, np.ndarray)
+            and len(values) > 0
+            and is_string_array(values, skipna=True)
+        ):
             result = result.astype(StringDtype(na_value=np.nan))
         return result
 
@@ -3369,6 +3375,7 @@ def read(
             if (
                 using_string_dtype()
                 and isinstance(values, np.ndarray)
+                and len(df) > 0
                 and is_string_array(values, skipna=True)
             ):
                 df = df.astype(StringDtype(na_value=np.nan))
@@ -4747,6 +4754,7 @@ def read(
             if (
                 using_string_dtype()
                 and isinstance(values, np.ndarray)
+                and len(df) > 0
                 and is_string_array(values, skipna=True)
             ):
                 df = df.astype(StringDtype(na_value=np.nan))

pandas/tests/io/pytables/test_round_trip.py

@@ -4,6 +4,8 @@
 import numpy as np
 import pytest
 
+from pandas._config import using_string_dtype
+
 from pandas._libs.tslibs import Timestamp
 from pandas.compat import is_platform_windows
 
@@ -66,6 +68,7 @@ def roundtrip(key, obj, **kwargs):
         tm.assert_frame_equal(result, expected)
 
 
+@pytest.mark.xfail(using_string_dtype(), reason="TODO(infer_string)")
 def test_long_strings(setup_path):
     # GH6166
     data = ["a" * 50] * 10
@@ -206,6 +209,7 @@ def test_put_integer(setup_path):
     _check_roundtrip(df, tm.assert_frame_equal, setup_path)
 
 
+@pytest.mark.xfail(using_string_dtype(), reason="TODO(infer_string)")
 def test_table_values_dtypes_roundtrip(setup_path):
     with ensure_clean_store(setup_path) as store:
         df1 = DataFrame({"a": [1, 2, 3]}, dtype="f8")
@@ -375,7 +379,7 @@ def test_timeseries_preepoch(setup_path, request):
 @pytest.mark.parametrize(
     "compression", [False, pytest.param(True, marks=td.skip_if_windows)]
 )
-def test_frame(compression, setup_path):
+def test_frame(compression, setup_path, using_infer_string):
     df = DataFrame(
         1.1 * np.arange(120).reshape((30, 4)),
         columns=Index(list("ABCD"), dtype=object),
@@ -386,20 +390,40 @@ def test_frame(compression, setup_path):
     df.iloc[0, 0] = np.nan
     df.iloc[5, 3] = np.nan
 
+    expected = df.copy()
+    if using_infer_string:
+        expected.index = expected.index.astype("str")
+        expected.columns = expected.columns.astype("str")
+
     _check_roundtrip_table(
-        df, tm.assert_frame_equal, path=setup_path, compression=compression
+        df,
+        tm.assert_frame_equal,
+        path=setup_path,
+        compression=compression,
+        expected=expected,
     )
     _check_roundtrip(
-        df, tm.assert_frame_equal, path=setup_path, compression=compression
+        df,
+        tm.assert_frame_equal,
+        path=setup_path,
+        compression=compression,
+        expected=expected,
     )
 
     tdf = DataFrame(
         np.random.default_rng(2).standard_normal((10, 4)),
         columns=Index(list("ABCD"), dtype=object),
         index=date_range("2000-01-01", periods=10, freq="B"),
     )
+    expected = tdf.copy()
+    if using_infer_string:
+        expected.columns = expected.columns.astype("str")
     _check_roundtrip(
-        tdf, tm.assert_frame_equal, path=setup_path, compression=compression
+        tdf,
+        tm.assert_frame_equal,
+        path=setup_path,
+        compression=compression,
+        expected=expected,
     )
 
     with ensure_clean_store(setup_path) as store:
@@ -410,7 +434,10 @@ def test_frame(compression, setup_path):
         assert recons._mgr.is_consolidated()
 
     # empty
-    _check_roundtrip(df[:0], tm.assert_frame_equal, path=setup_path)
+    expected = df[:0]
+    if using_infer_string:
+        expected.columns = expected.columns.astype("str")
+    _check_roundtrip(df[:0], tm.assert_frame_equal, path=setup_path, expected=expected)
 
 
 def test_empty_series_frame(setup_path):
@@ -442,9 +469,21 @@ def test_can_serialize_dates(setup_path):
     _check_roundtrip(frame, tm.assert_frame_equal, path=setup_path)
 
 
-def test_store_hierarchical(setup_path, multiindex_dataframe_random_data):
+def test_store_hierarchical(
+    setup_path, multiindex_dataframe_random_data, using_infer_string
+):
     frame = multiindex_dataframe_random_data
 
+    if using_infer_string:
+        msg = "Saving a MultiIndex with an extension dtype is not supported."
+        with pytest.raises(NotImplementedError, match=msg):
+            _check_roundtrip(frame, tm.assert_frame_equal, path=setup_path)
+        with pytest.raises(NotImplementedError, match=msg):
+            _check_roundtrip(frame.T, tm.assert_frame_equal, path=setup_path)
+        with pytest.raises(NotImplementedError, match=msg):
+            _check_roundtrip(frame["A"], tm.assert_series_equal, path=setup_path)
+        return
+
     _check_roundtrip(frame, tm.assert_frame_equal, path=setup_path)
     _check_roundtrip(frame.T, tm.assert_frame_equal, path=setup_path)
     _check_roundtrip(frame["A"], tm.assert_series_equal, path=setup_path)
@@ -459,7 +498,7 @@ def test_store_hierarchical(setup_path, multiindex_dataframe_random_data):
 @pytest.mark.parametrize(
     "compression", [False, pytest.param(True, marks=td.skip_if_windows)]
 )
-def test_store_mixed(compression, setup_path):
+def test_store_mixed(compression, setup_path, using_infer_string):
     def _make_one():
         df = DataFrame(
             1.1 * np.arange(120).reshape((30, 4)),
@@ -477,57 +516,91 @@ def _make_one():
     df1 = _make_one()
     df2 = _make_one()
 
-    _check_roundtrip(df1, tm.assert_frame_equal, path=setup_path)
-    _check_roundtrip(df2, tm.assert_frame_equal, path=setup_path)
+    expected = df1.copy()
+    if using_infer_string:
+        expected.index = expected.index.astype("str")
+        expected.columns = expected.columns.astype("str")
+    _check_roundtrip(df1, tm.assert_frame_equal, path=setup_path, expected=expected)
+
+    expected = df2.copy()
+    if using_infer_string:
+        expected.index = expected.index.astype("str")
+        expected.columns = expected.columns.astype("str")
+    _check_roundtrip(df2, tm.assert_frame_equal, path=setup_path, expected=expected)
 
     with ensure_clean_store(setup_path) as store:
         store["obj"] = df1
-        tm.assert_frame_equal(store["obj"], df1)
+        expected = df1.copy()
+        if using_infer_string:
+            expected.index = expected.index.astype("str")
+            expected.columns = expected.columns.astype("str")
+        tm.assert_frame_equal(store["obj"], expected)
+
         store["obj"] = df2
-        tm.assert_frame_equal(store["obj"], df2)
+        expected = df2.copy()
+        if using_infer_string:
+            expected.index = expected.index.astype("str")
+            expected.columns = expected.columns.astype("str")
+        tm.assert_frame_equal(store["obj"], expected)
 
     # check that can store Series of all of these types
+    expected = df1["obj1"]
+    if using_infer_string:
+        expected.index = expected.index.astype("str")
     _check_roundtrip(
         df1["obj1"],
         tm.assert_series_equal,
         path=setup_path,
         compression=compression,
+        expected=expected,
     )
+    expected = df1["bool1"]
+    if using_infer_string:
+        expected.index = expected.index.astype("str")
     _check_roundtrip(
         df1["bool1"],
         tm.assert_series_equal,
         path=setup_path,
         compression=compression,
+        expected=expected,
     )
+    expected = df1["int1"]
+    if using_infer_string:
+        expected.index = expected.index.astype("str")
     _check_roundtrip(
         df1["int1"],
         tm.assert_series_equal,
         path=setup_path,
         compression=compression,
+        expected=expected,
     )
 
 
-def _check_roundtrip(obj, comparator, path, compression=False, **kwargs):
+def _check_roundtrip(obj, comparator, path, compression=False, expected=None, **kwargs):
     options = {}
     if compression:
         options["complib"] = "blosc"
+    if expected is None:
+        expected = obj
 
     with ensure_clean_store(path, "w", **options) as store:
         store["obj"] = obj
         retrieved = store["obj"]
-        comparator(retrieved, obj, **kwargs)
+        comparator(retrieved, expected, **kwargs)
 
 
-def _check_roundtrip_table(obj, comparator, path, compression=False):
+def _check_roundtrip_table(obj, comparator, path, compression=False, expected=None):
     options = {}
     if compression:
         options["complib"] = "blosc"
+    if expected is None:
+        expected = obj
 
     with ensure_clean_store(path, "w", **options) as store:
         store.put("obj", obj, format="table")
         retrieved = store["obj"]
 
-        comparator(retrieved, obj)
+        comparator(retrieved, expected)
 
 
 def test_unicode_index(setup_path):
@@ -540,6 +613,7 @@ def test_unicode_index(setup_path):
     _check_roundtrip(s, tm.assert_series_equal, path=setup_path)
 
 
+@pytest.mark.xfail(using_string_dtype(), reason="TODO(infer_string)")
 def test_unicode_longer_encoded(setup_path):
     # GH 11234
     char = "\u0394"
@@ -565,6 +639,7 @@ def test_store_datetime_mixed(setup_path):
     _check_roundtrip(df, tm.assert_frame_equal, path=setup_path)
 
 
+@pytest.mark.xfail(using_string_dtype(), reason="TODO(infer_string)")
 def test_round_trip_equals(tmp_path, setup_path):
     # GH 9330
     df = DataFrame({"B": [1, 2], "A": ["x", "y"]})