parametrize pyarrow tests

Author: aaronchucarroll

pandas/tests/strings/test_get_dummies.py

@@ -4,7 +4,6 @@
 import pandas.util._test_decorators as td
 
 from pandas import (
-    ArrowDtype,
     DataFrame,
     Index,
     MultiIndex,
@@ -69,108 +68,26 @@ def test_get_dummies_with_dtype(any_string_dtype, dtype):
 
 
 @td.skip_if_no("pyarrow")
-def test_get_dummies_with_pyarrow_dtype_int8(any_string_dtype):
-    s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
-    result = s.str.get_dummies("|", dtype=ArrowDtype(pa.int8()))
-    expected = DataFrame(
-        [[1, 1, 0], [1, 0, 1], [0, 0, 0]],
-        columns=list("abc"),
-        dtype=ArrowDtype(pa.int8()),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-@td.skip_if_no("pyarrow")
-def test_get_dummies_with_pyarrow_dtype_uint8(any_string_dtype):
-    s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
-    result = s.str.get_dummies("|", dtype=ArrowDtype(pa.uint8()))
-    expected = DataFrame(
-        [[1, 1, 0], [1, 0, 1], [0, 0, 0]],
-        columns=list("abc"),
-        dtype=ArrowDtype(pa.uint8()),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-@td.skip_if_no("pyarrow")
-def test_get_dummies_with_pyarrow_dtype_int16(any_string_dtype):
-    s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
-    result = s.str.get_dummies("|", dtype=ArrowDtype(pa.int16()))
-    expected = DataFrame(
-        [[1, 1, 0], [1, 0, 1], [0, 0, 0]],
-        columns=list("abc"),
-        dtype=ArrowDtype(pa.int16()),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-@td.skip_if_no("pyarrow")
-def test_get_dummies_with_pyarrow_dtype_uint16(any_string_dtype):
-    s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
-    result = s.str.get_dummies("|", dtype=ArrowDtype(pa.uint16()))
-    expected = DataFrame(
-        [[1, 1, 0], [1, 0, 1], [0, 0, 0]],
-        columns=list("abc"),
-        dtype=ArrowDtype(pa.uint16()),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-@td.skip_if_no("pyarrow")
-def test_get_dummies_with_pyarrow_dtype_int32(any_string_dtype):
-    s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
-    result = s.str.get_dummies("|", dtype=ArrowDtype(pa.int32()))
-    expected = DataFrame(
-        [[1, 1, 0], [1, 0, 1], [0, 0, 0]],
-        columns=list("abc"),
-        dtype=ArrowDtype(pa.int32()),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-@td.skip_if_no("pyarrow")
-def test_get_dummies_with_pyarrow_dtype_uint32(any_string_dtype):
-    s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
-    result = s.str.get_dummies("|", dtype=ArrowDtype(pa.uint32()))
-    expected = DataFrame(
-        [[1, 1, 0], [1, 0, 1], [0, 0, 0]],
-        columns=list("abc"),
-        dtype=ArrowDtype(pa.uint32()),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-@td.skip_if_no("pyarrow")
-def test_get_dummies_with_pyarrow_dtype_int64(any_string_dtype):
-    s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
-    result = s.str.get_dummies("|", dtype=ArrowDtype(pa.int64()))
-    expected = DataFrame(
-        [[1, 1, 0], [1, 0, 1], [0, 0, 0]],
-        columns=list("abc"),
-        dtype=ArrowDtype(pa.int64()),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-@td.skip_if_no("pyarrow")
-def test_get_dummies_with_pyarrow_dtype_uint64(any_string_dtype):
-    s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
-    result = s.str.get_dummies("|", dtype=ArrowDtype(pa.uint64()))
-    expected = DataFrame(
-        [[1, 1, 0], [1, 0, 1], [0, 0, 0]],
-        columns=list("abc"),
-        dtype=ArrowDtype(pa.uint64()),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-@td.skip_if_no("pyarrow")
-def test_get_dummies_with_pyarrow_dtype_bool(any_string_dtype):
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        "int8[pyarrow]",
+        "uint8[pyarrow]",
+        "int16[pyarrow]",
+        "uint16[pyarrow]",
+        "int32[pyarrow]",
+        "uint32[pyarrow]",
+        "int64[pyarrow]",
+        "uint64[pyarrow]",
+        "bool[pyarrow]",
+    ],
+)
+def test_get_dummies_with_pyarrow_dtype(any_string_dtype, dtype):
     s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
-    result = s.str.get_dummies("|", dtype=ArrowDtype(pa.bool_()))
+    result = s.str.get_dummies("|", dtype=dtype)
     expected = DataFrame(
         [[1, 1, 0], [1, 0, 1], [0, 0, 0]],
         columns=list("abc"),
-        dtype=ArrowDtype(pa.bool_()),
+        dtype=dtype,
     )
     tm.assert_frame_equal(result, expected)