fix select_dtypes test + more test changes

Author: jorisvandenbossche

pandas/core/frame.py

@@ -4807,7 +4807,9 @@ def select_dtypes(self, include=None, exclude=None) -> DataFrame:
         -----
         * To select all *numeric* types, use ``np.number`` or ``'number'``
         * To select strings you must use the ``object`` dtype, but note that
-          this will return *all* object dtype columns
+          this will return *all* object dtype columns. With
+          ``pd.options.future.infer_string`` enabled, using ``"str"`` will
+          work to select all string columns.
         * See the `numpy dtype hierarchy
           <https://numpy.org/doc/stable/reference/arrays.scalars.html>`__
         * To select datetimes, use ``np.datetime64``, ``'datetime'`` or

pandas/tests/frame/methods/test_astype.py

@@ -186,7 +186,7 @@ def test_astype_str_float(self):
         tm.assert_frame_equal(result, expected)
 
     @pytest.mark.parametrize("dtype_class", [dict, Series])
-    def test_astype_dict_like(self, dtype_class, using_infer_string):
+    def test_astype_dict_like(self, dtype_class):
         # GH7271 & GH16717
         a = Series(date_range("2010-01-04", periods=5))
         b = Series(range(5))
@@ -201,10 +201,7 @@ def test_astype_dict_like(self, dtype_class, using_infer_string):
         expected = DataFrame(
             {
                 "a": a,
-                "b": Series(
-                    ["0", "1", "2", "3", "4"],
-                    dtype="str" if using_infer_string else "object",
-                ),
+                "b": Series(["0", "1", "2", "3", "4"], dtype="str"),
                 "c": c,
                 "d": Series([1.0, 2.0, 3.14, 4.0, 5.4], dtype="float32"),
             }

pandas/tests/frame/methods/test_reset_index.py

@@ -662,7 +662,7 @@ def test_reset_index_dtypes_on_empty_frame_with_multiindex(
     idx = MultiIndex.from_product([[0, 1], [0.5, 1.0], array])
     result = DataFrame(index=idx)[:0].reset_index().dtypes
     if using_infer_string and dtype == object:
-        dtype = "str"
+        dtype = pd.StringDtype(na_value=np.nan)
     expected = Series({"level_0": np.int64, "level_1": np.float64, "level_2": dtype})
     tm.assert_series_equal(result, expected)
 

pandas/tests/frame/methods/test_select_dtypes.py

@@ -1,8 +1,6 @@
 import numpy as np
 import pytest
 
-from pandas._config import using_string_dtype
-
 from pandas.core.dtypes.dtypes import ExtensionDtype
 
 import pandas as pd
@@ -52,7 +50,7 @@ def copy(self):
 
 
 class TestSelectDtypes:
-    def test_select_dtypes_include_using_list_like(self):
+    def test_select_dtypes_include_using_list_like(self, using_infer_string):
         df = DataFrame(
             {
                 "a": list("abc"),
@@ -96,6 +94,11 @@ def test_select_dtypes_include_using_list_like(self):
         with pytest.raises(NotImplementedError, match=r"^$"):
             df.select_dtypes(include=["period"])
 
+        if using_infer_string:
+            ri = df.select_dtypes(include=["str"])
+            ei = df[["a"]]
+            tm.assert_frame_equal(ri, ei)
+
     def test_select_dtypes_exclude_using_list_like(self):
         df = DataFrame(
             {
@@ -153,7 +156,7 @@ def test_select_dtypes_exclude_include_int(self, include):
         expected = df[["b", "c", "e"]]
         tm.assert_frame_equal(result, expected)
 
-    def test_select_dtypes_include_using_scalars(self):
+    def test_select_dtypes_include_using_scalars(self, using_infer_string):
         df = DataFrame(
             {
                 "a": list("abc"),
@@ -189,6 +192,11 @@ def test_select_dtypes_include_using_scalars(self):
         with pytest.raises(NotImplementedError, match=r"^$"):
             df.select_dtypes(include="period")
 
+        if using_infer_string:
+            ri = df.select_dtypes(include="str")
+            ei = df[["a"]]
+            tm.assert_frame_equal(ri, ei)
+
     def test_select_dtypes_exclude_using_scalars(self):
         df = DataFrame(
             {
@@ -347,10 +355,12 @@ def test_select_dtypes_datetime_with_tz(self):
         expected = df3.reindex(columns=[])
         tm.assert_frame_equal(result, expected)
 
-    @pytest.mark.xfail(using_string_dtype(), reason="TODO(infer_string)", strict=False)
     @pytest.mark.parametrize("dtype", [str, "str", np.bytes_, "S1", np.str_, "U1"])
     @pytest.mark.parametrize("arg", ["include", "exclude"])
-    def test_select_dtypes_str_raises(self, dtype, arg):
+    def test_select_dtypes_str_raises(self, dtype, arg, using_infer_string):
+        if using_infer_string and dtype == "str":
+            # this is tested below
+            pytest.skip("Selecting string columns works with future strings")
         df = DataFrame(
             {
                 "a": list("abc"),

pandas/tests/frame/test_block_internals.py

@@ -186,7 +186,9 @@ def test_construction_with_mixed(self, float_string_frame, using_infer_string):
         expected = Series(
             [np.dtype("float64")] * 4
             + [
-                np.dtype("object") if not using_infer_string else "str",
+                np.dtype("object")
+                if not using_infer_string
+                else pd.StringDtype(na_value=np.nan),
                 np.dtype("datetime64[us]"),
                 np.dtype("timedelta64[us]"),
             ],

pandas/tests/frame/test_stack_unstack.py

@@ -658,7 +658,11 @@ def test_unstack_dtypes(self, using_infer_string):
         df2["D"] = "foo"
         df3 = df2.unstack("B")
         result = df3.dtypes
-        dtype = "str" if using_infer_string else np.dtype("object")
+        dtype = (
+            pd.StringDtype(na_value=np.nan)
+            if using_infer_string
+            else np.dtype("object")
+        )
         expected = Series(
             [np.dtype("float64")] * 2 + [dtype] * 2,
             index=MultiIndex.from_arrays(

pandas/tests/groupby/test_apply.py

@@ -1013,7 +1013,7 @@ def test_groupby_apply_datetime_result_dtypes(using_infer_string):
     msg = "DataFrameGroupBy.apply operated on the grouping columns"
     with tm.assert_produces_warning(DeprecationWarning, match=msg):
         result = data.groupby("color").apply(lambda g: g.iloc[0]).dtypes
-    dtype = "str" if using_infer_string else object
+    dtype = pd.StringDtype(na_value=np.nan) if using_infer_string else object
     expected = Series(
         [np.dtype("datetime64[us]"), dtype, dtype, np.int64, dtype],
         index=["observation", "color", "mood", "intensity", "score"],

pandas/tests/groupby/test_categorical.py

@@ -134,8 +134,8 @@ def f(x):
         result = g.apply(f)
     expected = x.iloc[[0, 1]].copy()
     expected.index = Index([1, 2], name="person_id")
-    dtype = "str" if using_infer_string else object
-    expected["person_name"] = expected["person_name"].astype(dtype)
+    # dtype = "str" if using_infer_string else object
+    # expected["person_name"] = expected["person_name"].astype(dtype)
     tm.assert_frame_equal(result, expected)
 
 

pandas/tests/indexes/multi/test_constructors.py

@@ -850,7 +850,7 @@ 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 "str"
+    exp = "object" if not using_infer_string else pd.StringDtype(na_value=np.nan)
     expected = Series(
         ["int64", exp],
         index=MultiIndex.from_tuples([("a", "b"), ("c", "d")]),

pandas/tests/indexes/object/test_indexing.py

@@ -171,6 +171,7 @@ def test_get_indexer_non_unique_np_nats(self, np_nat_fixture, np_nat_fixture2):
 
 
 class TestSliceLocs:
+    # TODO(infer_string) parametrize over multiple string dtypes
     @pytest.mark.parametrize(
         "dtype",
         [
@@ -209,6 +210,7 @@ def test_slice_locs_negative_step(self, in_slice, expected, dtype):
         expected = Index(list(expected), dtype=dtype)
         tm.assert_index_equal(result, expected)
 
+    # TODO(infer_string) parametrize over multiple string dtypes
     @td.skip_if_no("pyarrow")
     def test_slice_locs_negative_step_oob(self):
         index = Index(list("bcdxy"), dtype="string[pyarrow_numpy]")