Adjust Series specific tests for string option (#55538)

Author: phofl

pandas/conftest.py

@@ -2024,6 +2024,14 @@ def warn_copy_on_write() -> bool:
  )
 
 
+@pytest.fixture
+def using_infer_string() -> bool:
+ """
+ Fixture to check if infer_string is enabled.
+ """
+ return pd.options.future.infer_string
+
+
 warsaws = ["Europe/Warsaw", "dateutil/Europe/Warsaw"]
 if zoneinfo is not None:
  warsaws.append(zoneinfo.ZoneInfo("Europe/Warsaw")) # type: ignore[arg-type]

pandas/tests/series/accessors/test_dt_accessor.py

@@ -586,13 +586,15 @@ def test_strftime_dt64_days(self):
  # dtype may be S10 or U10 depending on python version
  tm.assert_index_equal(result, expected)
 
- def test_strftime_period_days(self):
+ def test_strftime_period_days(self, using_infer_string):
  period_index = period_range("20150301", periods=5)
  result = period_index.strftime("%Y/%m/%d")
  expected = Index(
  ["2015/03/01", "2015/03/02", "2015/03/03", "2015/03/04", "2015/03/05"],
  dtype="=U10",
  )
+ if using_infer_string:
+ expected = expected.astype("string[pyarrow_numpy]")
  tm.assert_index_equal(result, expected)
 
  def test_strftime_dt64_microsecond_resolution(self):

pandas/tests/series/indexing/test_delitem.py

@@ -31,19 +31,16 @@ def test_delitem(self):
  del s[0]
  tm.assert_series_equal(s, Series(dtype="int64", index=Index([], dtype="int64")))
 
- def test_delitem_object_index(self):
+ def test_delitem_object_index(self, using_infer_string):
  # Index(dtype=object)
- s = Series(1, index=["a"])
+ dtype = "string[pyarrow_numpy]" if using_infer_string else object
+ s = Series(1, index=Index(["a"], dtype=dtype))
  del s["a"]
- tm.assert_series_equal(
- s, Series(dtype="int64", index=Index([], dtype="object"))
- )
+ tm.assert_series_equal(s, Series(dtype="int64", index=Index([], dtype=dtype)))
  s["a"] = 1
- tm.assert_series_equal(s, Series(1, index=["a"]))
+ tm.assert_series_equal(s, Series(1, index=Index(["a"], dtype=dtype)))
  del s["a"]
- tm.assert_series_equal(
- s, Series(dtype="int64", index=Index([], dtype="object"))
- )
+ tm.assert_series_equal(s, Series(dtype="int64", index=Index([], dtype=dtype)))
 
  def test_delitem_missing_key(self):
  # empty

pandas/tests/series/indexing/test_getitem.py

@@ -71,7 +71,7 @@ def test_getitem_unrecognized_scalar(self):
  def test_getitem_negative_out_of_bounds(self):
  ser = Series(["a"] * 10, index=["a"] * 10)
 
- msg = "index -11 is out of bounds for axis 0 with size 10"
+ msg = "index -11 is out of bounds for axis 0 with size 10|index out of bounds"
  warn_msg = "Series.__getitem__ treating keys as positions is deprecated"
  with pytest.raises(IndexError, match=msg):
  with tm.assert_produces_warning(FutureWarning, match=warn_msg):
@@ -363,7 +363,9 @@ def test_getitem_no_matches(self, box):
  key = Series(["C"], dtype=object)
  key = box(key)
 
- msg = r"None of \[Index\(\['C'\], dtype='object'\)\] are in the \[index\]"
+ msg = (
+ r"None of \[Index\(\['C'\], dtype='object|string'\)\] are in the \[index\]"
+ )
  with pytest.raises(KeyError, match=msg):
  ser[key]
 
@@ -437,7 +439,7 @@ def test_getitem_boolean_empty(self):
 
  # GH#5877
  # indexing with empty series
- ser = Series(["A", "B"])
+ ser = Series(["A", "B"], dtype=object)
  expected = Series(dtype=object, index=Index([], dtype="int64"))
  result = ser[Series([], dtype=object)]
  tm.assert_series_equal(result, expected)

pandas/tests/series/indexing/test_setitem.py

@@ -2,6 +2,7 @@
  date,
  datetime,
 )
+from decimal import Decimal
 
 import numpy as np
 import pytest
@@ -175,7 +176,8 @@ class TestSetitemScalarIndexer:
  def test_setitem_negative_out_of_bounds(self):
  ser = Series(["a"] * 10, index=["a"] * 10)
 
- msg = "index -11 is out of bounds for axis 0 with size 10"
+ # string index falls back to positional
+ msg = "index -11|-1 is out of bounds for axis 0 with size 10"
  warn_msg = "Series.__setitem__ treating keys as positions is deprecated"
  with pytest.raises(IndexError, match=msg):
  with tm.assert_produces_warning(FutureWarning, match=warn_msg):
@@ -527,8 +529,12 @@ def test_setitem_empty_series_timestamp_preserves_dtype(self):
  Timedelta("9 days").to_pytimedelta(),
  ],
  )
- def test_append_timedelta_does_not_cast(self, td):
+ def test_append_timedelta_does_not_cast(self, td, using_infer_string, request):
  # GH#22717 inserting a Timedelta should _not_ cast to int64
+ if using_infer_string and not isinstance(td, Timedelta):
+ # TODO: GH#56010
+ request.applymarker(pytest.mark.xfail(reason="inferred as string"))
+
  expected = Series(["x", td], index=[0, "td"], dtype=object)
 
  ser = Series(["x"])
@@ -595,13 +601,21 @@ def test_setitem_enlarge_with_na(
  expected = Series(expected_values, dtype=target_dtype)
  tm.assert_series_equal(ser, expected)
 
- def test_setitem_enlargement_object_none(self, nulls_fixture):
+ def test_setitem_enlargement_object_none(self, nulls_fixture, using_infer_string):
  # GH#48665
  ser = Series(["a", "b"])
  ser[3] = nulls_fixture
- expected = Series(["a", "b", nulls_fixture], index=[0, 1, 3])
+ dtype = (
+ "string[pyarrow_numpy]"
+ if using_infer_string and not isinstance(nulls_fixture, Decimal)
+ else object
+ )
+ expected = Series(["a", "b", nulls_fixture], index=[0, 1, 3], dtype=dtype)
  tm.assert_series_equal(ser, expected)
- assert ser[3] is nulls_fixture
+ if using_infer_string:
+ ser[3] is np.nan
+ else:
+ assert ser[3] is nulls_fixture
 
 
 def test_setitem_scalar_into_readonly_backing_data():
@@ -845,20 +859,28 @@ def test_series_where(self, obj, key, expected, warn, val, is_inplace):
 
  self._check_inplace(is_inplace, orig, arr, obj)
 
- def test_index_where(self, obj, key, expected, warn, val):
+ def test_index_where(self, obj, key, expected, warn, val, using_infer_string):
  mask = np.zeros(obj.shape, dtype=bool)
  mask[key] = True
 
- res = Index(obj).where(~mask, val)
- expected_idx = Index(expected, dtype=expected.dtype)
- tm.assert_index_equal(res, expected_idx)
+ if using_infer_string and obj.dtype == object:
+ with pytest.raises(TypeError, match="Scalar must"):
+ Index(obj).where(~mask, val)
+ else:
+ res = Index(obj).where(~mask, val)
+ expected_idx = Index(expected, dtype=expected.dtype)
+ tm.assert_index_equal(res, expected_idx)
 
- def test_index_putmask(self, obj, key, expected, warn, val):
+ def test_index_putmask(self, obj, key, expected, warn, val, using_infer_string):
  mask = np.zeros(obj.shape, dtype=bool)
  mask[key] = True
 
- res = Index(obj).putmask(mask, val)
- tm.assert_index_equal(res, Index(expected, dtype=expected.dtype))
+ if using_infer_string and obj.dtype == object:
+ with pytest.raises(TypeError, match="Scalar must"):
+ Index(obj).putmask(mask, val)
+ else:
+ res = Index(obj).putmask(mask, val)
+ tm.assert_index_equal(res, Index(expected, dtype=expected.dtype))
 
 
 @pytest.mark.parametrize(

pandas/tests/series/indexing/test_where.py

@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 
+from pandas._config import using_pyarrow_string_dtype
+
 from pandas.core.dtypes.common import is_integer
 
 import pandas as pd
@@ -230,6 +232,7 @@ def test_where_ndframe_align():
  tm.assert_series_equal(out, expected)
 
 
+@pytest.mark.xfail(using_pyarrow_string_dtype(), reason="can't set ints into string")
 def test_where_setitem_invalid():
  # GH 2702
  # make sure correct exceptions are raised on invalid list assignment

pandas/tests/series/methods/test_astype.py

@@ -76,7 +76,7 @@ def test_astype_dict_like(self, dtype_class):
 
  dt1 = dtype_class({"abc": str})
  result = ser.astype(dt1)
- expected = Series(["0", "2", "4", "6", "8"], name="abc")
+ expected = Series(["0", "2", "4", "6", "8"], name="abc", dtype=object)
  tm.assert_series_equal(result, expected)
 
  dt2 = dtype_class({"abc": "float64"})
@@ -170,10 +170,12 @@ def test_astype_empty_constructor_equality(self, dtype):
  Series([string.digits * 10, rand_str(63), rand_str(64), np.nan, 1.0]),
  ],
  )
- def test_astype_str_map(self, dtype, series):
+ def test_astype_str_map(self, dtype, series, using_infer_string):
  # see GH#4405
  result = series.astype(dtype)
  expected = series.map(str)
+ if using_infer_string:
+ expected = expected.astype(object)
  tm.assert_series_equal(result, expected)
 
  def test_astype_float_to_period(self):
@@ -283,13 +285,13 @@ def test_astype_str_cast_dt64(self):
  ts = Series([Timestamp("2010-01-04 00:00:00")])
  res = ts.astype(str)
 
- expected = Series(["2010-01-04"])
+ expected = Series(["2010-01-04"], dtype=object)
  tm.assert_series_equal(res, expected)
 
  ts = Series([Timestamp("2010-01-04 00:00:00", tz="US/Eastern")])
  res = ts.astype(str)
 
- expected = Series(["2010-01-04 00:00:00-05:00"])
+ expected = Series(["2010-01-04 00:00:00-05:00"], dtype=object)
  tm.assert_series_equal(res, expected)
 
  def test_astype_str_cast_td64(self):
@@ -298,7 +300,7 @@ def test_astype_str_cast_td64(self):
  td = Series([Timedelta(1, unit="d")])
  ser = td.astype(str)
 
- expected = Series(["1 days"])
+ expected = Series(["1 days"], dtype=object)
  tm.assert_series_equal(ser, expected)
 
  def test_dt64_series_astype_object(self):
@@ -345,7 +347,7 @@ def test_astype_from_float_to_str(self, dtype):
  # https://github.com/pandas-dev/pandas/issues/36451
  ser = Series([0.1], dtype=dtype)
  result = ser.astype(str)
- expected = Series(["0.1"])
+ expected = Series(["0.1"], dtype=object)
  tm.assert_series_equal(result, expected)
 
  @pytest.mark.parametrize(
@@ -416,7 +418,7 @@ def test_astype_cast_object_int(self):
 
  tm.assert_series_equal(result, Series(np.arange(1, 5)))
 
- def test_astype_unicode(self):
+ def test_astype_unicode(self, using_infer_string):
  # see GH#7758: A bit of magic is required to set
  # default encoding to utf-8
  digits = string.digits
@@ -433,12 +435,14 @@ def test_astype_unicode(self):
  item = "野菜食べないとやばい"
  ser = Series([item.encode()])
  result = ser.astype(np.str_)
- expected = Series([item])
+ expected = Series([item], dtype=object)
  tm.assert_series_equal(result, expected)
 
  for ser in test_series:
  res = ser.astype(np.str_)
  expec = ser.map(str)
+ if using_infer_string:
+ expec = expec.astype(object)
  tm.assert_series_equal(res, expec)
 
  # Restore the former encoding
@@ -534,12 +538,12 @@ def test_astype_categorical_to_other(self):
  expected = ser
  tm.assert_series_equal(ser.astype("category"), expected)
  tm.assert_series_equal(ser.astype(CategoricalDtype()), expected)
- msg = r"Cannot cast object dtype to float64"
+ msg = r"Cannot cast object|string dtype to float64"
  with pytest.raises(ValueError, match=msg):
  ser.astype("float64")
 
  cat = Series(Categorical(["a", "b", "b", "a", "a", "c", "c", "c"]))
- exp = Series(["a", "b", "b", "a", "a", "c", "c", "c"])
+ exp = Series(["a", "b", "b", "a", "a", "c", "c", "c"], dtype=object)
  tm.assert_series_equal(cat.astype("str"), exp)
  s2 = Series(Categorical(["1", "2", "3", "4"]))
  exp2 = Series([1, 2, 3, 4]).astype("int")

pandas/tests/series/methods/test_combine_first.py

@@ -53,7 +53,7 @@ def test_combine_first(self):
  # mixed types
  index = tm.makeStringIndex(20)
  floats = Series(np.random.default_rng(2).standard_normal(20), index=index)
- strings = Series(tm.makeStringIndex(10), index=index[::2])
+ strings = Series(tm.makeStringIndex(10), index=index[::2], dtype=object)
 
  combined = strings.combine_first(floats)
 

pandas/tests/series/methods/test_convert_dtypes.py

@@ -186,6 +186,7 @@ def test_convert_dtypes(
  self,
  test_cases,
  params,
+ using_infer_string,
  ):
  data, maindtype, expected_default, expected_other = test_cases
  if (
@@ -219,6 +220,16 @@ def test_convert_dtypes(
  for spec, dtype in expected_other.items():
  if all(params_dict[key] is val for key, val in zip(spec[::2], spec[1::2])):
  expected_dtype = dtype
+ if (
+ using_infer_string
+ and expected_default == "string"
+ and expected_dtype == object
+ and params[0]
+ and not params[1]
+ ):
+ # If we would convert with convert strings then infer_objects converts
+ # with the option
+ expected_dtype = "string[pyarrow_numpy]"
 
  expected = pd.Series(data, dtype=expected_dtype)
  tm.assert_series_equal(result, expected)