Adjust Index specific tests for string option (#56074)

phofl · web-flow · commit 91ddc8b66a0e · 2023-12-09T11:20:50.000-08:00
* BUG: setitem casting object Index to arrow strings

* Fix

* Start fixing index tests

* BUG: Index.isin raising for arrow strings and null set

* Fix more tests

* TST: Fix shares_memory for arrow string dtype

* TST: Fix shares_memory for arrow string dtype

* TST: Fix shares_memory for arrow string dtype

* Fix more tests

* BUG: Index.getitem returning wrong result with negative step for arrow

* Update

* Update

* Fix

* Update array.py

* Fix

* Move

* Move

* Fix

* Add gh ref

* Update v2.1.4.rst

* Finish

* Update

* Update test_base.py

* Update test_old_base.py

* Update conftest.py

* Update conftest.py

* Update test_old_base.py

* Update

* Update test_setops.py

* Fix pre-commit
diff --git a/pandas/conftest.py b/pandas/conftest.py
@@ -1903,7 +1903,7 @@ def using_copy_on_write() -> bool:
 @pytest.fixture
 def warn_copy_on_write() -> bool:
     """
-    Fixture to check if Copy-on-Write is enabled.
+    Fixture to check if Copy-on-Write is in warning mode.
     """
     return (
         pd.options.mode.copy_on_write == "warn"
@@ -1914,9 +1914,9 @@ def warn_copy_on_write() -> bool:
 @pytest.fixture
 def using_infer_string() -> bool:
     """
-    Fixture to check if infer_string is enabled.
+    Fixture to check if infer string option is enabled.
     """
-    return pd.options.future.infer_string
+    return pd.options.future.infer_string is True
 
 
 warsaws = ["Europe/Warsaw", "dateutil/Europe/Warsaw"]
diff --git a/pandas/tests/indexes/base_class/test_formats.py b/pandas/tests/indexes/base_class/test_formats.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+from pandas._config import using_pyarrow_string_dtype
 import pandas._config.config as cf
 
 from pandas import Index
@@ -15,6 +16,7 @@ def test_repr_is_valid_construction_code(self):
         res = eval(repr(idx))
         tm.assert_index_equal(res, idx)
 
+    @pytest.mark.xfail(using_pyarrow_string_dtype(), reason="repr different")
     @pytest.mark.parametrize(
         "index,expected",
         [
@@ -79,6 +81,7 @@ def test_string_index_repr(self, index, expected):
         result = repr(index)
         assert result == expected
 
+    @pytest.mark.xfail(using_pyarrow_string_dtype(), reason="repr different")
     @pytest.mark.parametrize(
         "index,expected",
         [
diff --git a/pandas/tests/indexes/base_class/test_reshape.py b/pandas/tests/indexes/base_class/test_reshape.py
@@ -33,13 +33,15 @@ def test_insert(self):
 
         # test empty
         null_index = Index([])
-        tm.assert_index_equal(Index(["a"]), null_index.insert(0, "a"))
+        tm.assert_index_equal(Index(["a"], dtype=object), null_index.insert(0, "a"))
 
-    def test_insert_missing(self, nulls_fixture):
+    def test_insert_missing(self, nulls_fixture, using_infer_string):
         # GH#22295
         # test there is no mangling of NA values
-        expected = Index(["a", nulls_fixture, "b", "c"])
-        result = Index(list("abc")).insert(1, nulls_fixture)
+        expected = Index(["a", nulls_fixture, "b", "c"], dtype=object)
+        result = Index(list("abc"), dtype=object).insert(
+            1, Index([nulls_fixture], dtype=object)
+        )
         tm.assert_index_equal(result, expected)
 
     @pytest.mark.parametrize(
diff --git a/pandas/tests/indexes/base_class/test_setops.py b/pandas/tests/indexes/base_class/test_setops.py
@@ -154,7 +154,7 @@ def test_intersection_str_dates(self, sort):
     def test_intersection_non_monotonic_non_unique(self, index2, expected_arr, sort):
         # non-monotonic non-unique
         index1 = Index(["A", "B", "A", "C"])
-        expected = Index(expected_arr, dtype="object")
+        expected = Index(expected_arr)
         result = index1.intersection(index2, sort=sort)
         if sort is None:
             expected = expected.sort_values()
diff --git a/pandas/tests/indexes/categorical/test_astype.py b/pandas/tests/indexes/categorical/test_astype.py
@@ -18,7 +18,7 @@ def test_astype(self):
         ci = CategoricalIndex(list("aabbca"), categories=list("cab"), ordered=False)
 
         result = ci.astype(object)
-        tm.assert_index_equal(result, Index(np.array(ci)))
+        tm.assert_index_equal(result, Index(np.array(ci), dtype=object))
 
         # this IS equal, but not the same class
         assert result.equals(ci)
diff --git a/pandas/tests/indexes/categorical/test_category.py b/pandas/tests/indexes/categorical/test_category.py
@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 
+from pandas._config import using_pyarrow_string_dtype
+
 from pandas._libs import index as libindex
 from pandas._libs.arrays import NDArrayBacked
 
@@ -47,7 +49,7 @@ def test_insert(self, simple_index):
 
         # invalid -> cast to object
         expected = ci.astype(object).insert(0, "d")
-        result = ci.insert(0, "d")
+        result = ci.insert(0, "d").astype(object)
         tm.assert_index_equal(result, expected, exact=True)
 
         # GH 18295 (test missing)
@@ -194,6 +196,7 @@ def test_unique(self, data, categories, expected_data, ordered):
         expected = CategoricalIndex(expected_data, dtype=dtype)
         tm.assert_index_equal(idx.unique(), expected)
 
+    @pytest.mark.xfail(using_pyarrow_string_dtype(), reason="repr doesn't roundtrip")
     def test_repr_roundtrip(self):
         ci = CategoricalIndex(["a", "b"], categories=["a", "b"], ordered=True)
         str(ci)
diff --git a/pandas/tests/indexes/categorical/test_formats.py b/pandas/tests/indexes/categorical/test_formats.py
@@ -1,6 +1,9 @@
 """
 Tests for CategoricalIndex.__repr__ and related methods.
 """
+import pytest
+
+from pandas._config import using_pyarrow_string_dtype
 import pandas._config.config as cf
 
 from pandas import CategoricalIndex
@@ -16,6 +19,7 @@ def test_format_different_scalar_lengths(self):
         with tm.assert_produces_warning(FutureWarning, match=msg):
             assert idx.format() == expected
 
+    @pytest.mark.xfail(using_pyarrow_string_dtype(), reason="repr different")
     def test_string_categorical_index_repr(self):
         # short
         idx = CategoricalIndex(["a", "bb", "ccc"])
diff --git a/pandas/tests/indexes/categorical/test_reindex.py b/pandas/tests/indexes/categorical/test_reindex.py
@@ -40,7 +40,7 @@ def test_reindex_duplicate_target(self):
         # See GH25459
         cat = CategoricalIndex(["a", "b", "c"], categories=["a", "b", "c", "d"])
         res, indexer = cat.reindex(["a", "c", "c"])
-        exp = Index(["a", "c", "c"], dtype="object")
+        exp = Index(["a", "c", "c"])
         tm.assert_index_equal(res, exp, exact=True)
         tm.assert_numpy_array_equal(indexer, np.array([0, 2, 2], dtype=np.intp))
 
diff --git a/pandas/tests/indexes/datetimes/methods/test_map.py b/pandas/tests/indexes/datetimes/methods/test_map.py
@@ -16,7 +16,7 @@ def test_map(self):
 
         f = lambda x: x.strftime("%Y%m%d")
         result = rng.map(f)
-        exp = Index([f(x) for x in rng], dtype="<U8")
+        exp = Index([f(x) for x in rng])
         tm.assert_index_equal(result, exp)
 
     def test_map_fallthrough(self, capsys):
diff --git a/pandas/tests/indexes/interval/test_formats.py b/pandas/tests/indexes/interval/test_formats.py
@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 
+from pandas._config import using_pyarrow_string_dtype
+
 from pandas import (
     DataFrame,
     DatetimeIndex,
@@ -31,13 +33,16 @@ class TestIntervalIndexRendering:
             (DataFrame, ("            0\n(0.0, 1.0]  a\nNaN         b\n(2.0, 3.0]  c")),
         ],
     )
-    def test_repr_missing(self, constructor, expected):
+    def test_repr_missing(self, constructor, expected, using_infer_string, request):
         # GH 25984
+        if using_infer_string and constructor is Series:
+            request.applymarker(pytest.mark.xfail(reason="repr different"))
         index = IntervalIndex.from_tuples([(0, 1), np.nan, (2, 3)])
         obj = constructor(list("abc"), index=index)
         result = repr(obj)
         assert result == expected
 
+    @pytest.mark.xfail(using_pyarrow_string_dtype(), reason="repr different")
     def test_repr_floats(self):
         # GH 32553
 
diff --git a/pandas/tests/indexes/multi/test_constructors.py b/pandas/tests/indexes/multi/test_constructors.py
@@ -847,11 +847,14 @@ def test_multiindex_inference_consistency():
     assert lev.dtype == object
 
 
-def test_dtype_representation():
+def test_dtype_representation(using_infer_string):
     # GH#46900
     pmidx = MultiIndex.from_arrays([[1], ["a"]], names=[("a", "b"), ("c", "d")])
     result = pmidx.dtypes
+    exp = "object" if not using_infer_string else "string"
     expected = Series(
-        ["int64", "object"], index=MultiIndex.from_tuples([("a", "b"), ("c", "d")])
+        ["int64", exp],
+        index=MultiIndex.from_tuples([("a", "b"), ("c", "d")]),
+        dtype=object,
     )
     tm.assert_series_equal(result, expected)
diff --git a/pandas/tests/indexes/multi/test_get_set.py b/pandas/tests/indexes/multi/test_get_set.py
@@ -34,23 +34,25 @@ def test_get_level_number_integer(idx):
         idx._get_level_number("fourth")
 
 
-def test_get_dtypes():
+def test_get_dtypes(using_infer_string):
     # Test MultiIndex.dtypes (# Gh37062)
     idx_multitype = MultiIndex.from_product(
         [[1, 2, 3], ["a", "b", "c"], pd.date_range("20200101", periods=2, tz="UTC")],
         names=["int", "string", "dt"],
     )
+
+    exp = "object" if not using_infer_string else "string"
     expected = pd.Series(
         {
             "int": np.dtype("int64"),
-            "string": np.dtype("O"),
+            "string": exp,
             "dt": DatetimeTZDtype(tz="utc"),
         }
     )
     tm.assert_series_equal(expected, idx_multitype.dtypes)
 
 
-def test_get_dtypes_no_level_name():
+def test_get_dtypes_no_level_name(using_infer_string):
     # Test MultiIndex.dtypes (# GH38580 )
     idx_multitype = MultiIndex.from_product(
         [
@@ -59,17 +61,18 @@ def test_get_dtypes_no_level_name():
             pd.date_range("20200101", periods=2, tz="UTC"),
         ],
     )
+    exp = "object" if not using_infer_string else "string"
     expected = pd.Series(
         {
             "level_0": np.dtype("int64"),
-            "level_1": np.dtype("O"),
+            "level_1": exp,
             "level_2": DatetimeTZDtype(tz="utc"),
         }
     )
     tm.assert_series_equal(expected, idx_multitype.dtypes)
 
 
-def test_get_dtypes_duplicate_level_names():
+def test_get_dtypes_duplicate_level_names(using_infer_string):
     # Test MultiIndex.dtypes with non-unique level names (# GH45174)
     result = MultiIndex.from_product(
         [
@@ -79,8 +82,9 @@ def test_get_dtypes_duplicate_level_names():
         ],
         names=["A", "A", "A"],
     ).dtypes
+    exp = "object" if not using_infer_string else "string"
     expected = pd.Series(
-        [np.dtype("int64"), np.dtype("O"), DatetimeTZDtype(tz="utc")],
+        [np.dtype("int64"), exp, DatetimeTZDtype(tz="utc")],
         index=["A", "A", "A"],
     )
     tm.assert_series_equal(result, expected)
diff --git a/pandas/tests/indexes/multi/test_reindex.py b/pandas/tests/indexes/multi/test_reindex.py
@@ -75,11 +75,14 @@ def test_reindex_lvl_preserves_names_when_target_is_list_or_array():
     assert idx.reindex([], level=1)[0].names == ["foo", "bar"]
 
 
-def test_reindex_lvl_preserves_type_if_target_is_empty_list_or_array():
+def test_reindex_lvl_preserves_type_if_target_is_empty_list_or_array(
+    using_infer_string,
+):
     # GH7774
     idx = MultiIndex.from_product([[0, 1], ["a", "b"]])
     assert idx.reindex([], level=0)[0].levels[0].dtype.type == np.int64
-    assert idx.reindex([], level=1)[0].levels[1].dtype.type == np.object_
+    exp = np.object_ if not using_infer_string else str
+    assert idx.reindex([], level=1)[0].levels[1].dtype.type == exp
 
     # case with EA levels
     cat = pd.Categorical(["foo", "bar"])
diff --git a/pandas/tests/indexes/multi/test_setops.py b/pandas/tests/indexes/multi/test_setops.py
@@ -263,19 +263,23 @@ def test_union(idx, sort):
         assert result.equals(idx)
 
 
-def test_union_with_regular_index(idx):
+def test_union_with_regular_index(idx, using_infer_string):
     other = Index(["A", "B", "C"])
 
     result = other.union(idx)
     assert ("foo", "one") in result
     assert "B" in result
 
-    msg = "The values in the array are unorderable"
-    with tm.assert_produces_warning(RuntimeWarning, match=msg):
-        result2 = idx.union(other)
-    # This is more consistent now, if sorting fails then we don't sort at all
-    # in the MultiIndex case.
-    assert not result.equals(result2)
+    if using_infer_string:
+        with pytest.raises(NotImplementedError, match="Can only union"):
+            idx.union(other)
+    else:
+        msg = "The values in the array are unorderable"
+        with tm.assert_produces_warning(RuntimeWarning, match=msg):
+            result2 = idx.union(other)
+        # This is more consistent now, if sorting fails then we don't sort at all
+        # in the MultiIndex case.
+        assert not result.equals(result2)
 
 
 def test_intersection(idx, sort):
@@ -756,7 +760,12 @@ def test_intersection_keep_ea_dtypes(val, any_numeric_ea_dtype):
 
 def test_union_with_na_when_constructing_dataframe():
     # GH43222
-    series1 = Series((1,), index=MultiIndex.from_tuples(((None, None),)))
+    series1 = Series(
+        (1,),
+        index=MultiIndex.from_arrays(
+            [Series([None], dtype="string"), Series([None], dtype="string")]
+        ),
+    )
     series2 = Series((10, 20), index=MultiIndex.from_tuples(((None, None), ("a", "b"))))
     result = DataFrame([series1, series2])
     expected = DataFrame({(np.nan, np.nan): [1.0, 10.0], ("a", "b"): [np.nan, 20.0]})
diff --git a/pandas/tests/indexes/object/test_astype.py b/pandas/tests/indexes/object/test_astype.py
@@ -20,7 +20,7 @@ def test_astype_str_from_bytes():
 
     # while we're here, check that Series.astype behaves the same
     result = Series(idx).astype(str)
-    expected = Series(expected)
+    expected = Series(expected, dtype=object)
     tm.assert_series_equal(result, expected)
 
 
diff --git a/pandas/tests/indexes/object/test_indexing.py b/pandas/tests/indexes/object/test_indexing.py
diff --git a/pandas/tests/indexes/period/methods/test_astype.py b/pandas/tests/indexes/period/methods/test_astype.py
diff --git a/pandas/tests/indexes/test_base.py b/pandas/tests/indexes/test_base.py
diff --git a/pandas/tests/indexes/test_old_base.py b/pandas/tests/indexes/test_old_base.py
diff --git a/pandas/tests/indexes/test_setops.py b/pandas/tests/indexes/test_setops.py
diff --git a/pandas/tests/indexes/timedeltas/methods/test_astype.py b/pandas/tests/indexes/timedeltas/methods/test_astype.py
