CLN: Split up Boolean array tests (#32780)

Author: dsaxton

pandas/tests/arrays/boolean/__init__.py

@@ -0,0 +1,42 @@
+import numpy as np
+import pytest
+
+import pandas as pd
+from pandas.tests.extension.base import BaseOpsUtil
+
+
+@pytest.fixture
+def data():
+    return pd.array(
+        [True, False] * 4 + [np.nan] + [True, False] * 44 + [np.nan] + [True, False],
+        dtype="boolean",
+    )
+
+
+class TestArithmeticOps(BaseOpsUtil):
+    def test_error(self, data, all_arithmetic_operators):
+        # invalid ops
+
+        op = all_arithmetic_operators
+        s = pd.Series(data)
+        ops = getattr(s, op)
+        opa = getattr(data, op)
+
+        # invalid scalars
+        with pytest.raises(TypeError):
+            ops("foo")
+        with pytest.raises(TypeError):
+            ops(pd.Timestamp("20180101"))
+
+        # invalid array-likes
+        if op not in ("__mul__", "__rmul__"):
+            # TODO(extension) numpy's mul with object array sees booleans as numbers
+            with pytest.raises(TypeError):
+                ops(pd.Series("foo", index=s.index))
+
+        # 2d
+        result = opa(pd.DataFrame({"A": s}))
+        assert result is NotImplemented
+
+        with pytest.raises(NotImplementedError):
+            opa(np.arange(len(s)).reshape(-1, len(s)))

pandas/tests/arrays/boolean/test_arithmetic.py

@@ -0,0 +1,53 @@
+import numpy as np
+import pytest
+
+import pandas as pd
+import pandas._testing as tm
+
+
+def test_astype():
+    # with missing values
+    arr = pd.array([True, False, None], dtype="boolean")
+
+    with pytest.raises(ValueError, match="cannot convert NA to integer"):
+        arr.astype("int64")
+
+    with pytest.raises(ValueError, match="cannot convert float NaN to"):
+        arr.astype("bool")
+
+    result = arr.astype("float64")
+    expected = np.array([1, 0, np.nan], dtype="float64")
+    tm.assert_numpy_array_equal(result, expected)
+
+    result = arr.astype("str")
+    expected = np.array(["True", "False", "<NA>"], dtype="object")
+    tm.assert_numpy_array_equal(result, expected)
+
+    # no missing values
+    arr = pd.array([True, False, True], dtype="boolean")
+    result = arr.astype("int64")
+    expected = np.array([1, 0, 1], dtype="int64")
+    tm.assert_numpy_array_equal(result, expected)
+
+    result = arr.astype("bool")
+    expected = np.array([True, False, True], dtype="bool")
+    tm.assert_numpy_array_equal(result, expected)
+
+
+def test_astype_to_boolean_array():
+    # astype to BooleanArray
+    arr = pd.array([True, False, None], dtype="boolean")
+
+    result = arr.astype("boolean")
+    tm.assert_extension_array_equal(result, arr)
+    result = arr.astype(pd.BooleanDtype())
+    tm.assert_extension_array_equal(result, arr)
+
+
+def test_astype_to_integer_array():
+    # astype to IntegerArray
+    arr = pd.array([True, False, None], dtype="boolean")
+
+    result = arr.astype("Int64")
+    expected = pd.array([1, 0, None], dtype="Int64")
+    tm.assert_extension_array_equal(result, expected)

pandas/tests/arrays/boolean/test_astype.py

@@ -0,0 +1,94 @@
+import numpy as np
+import pytest
+
+import pandas as pd
+import pandas._testing as tm
+from pandas.arrays import BooleanArray
+from pandas.tests.extension.base import BaseOpsUtil
+
+
+@pytest.fixture
+def data():
+    return pd.array(
+        [True, False] * 4 + [np.nan] + [True, False] * 44 + [np.nan] + [True, False],
+        dtype="boolean",
+    )
+
+
+class TestComparisonOps(BaseOpsUtil):
+    def _compare_other(self, data, op_name, other):
+        op = self.get_op_from_name(op_name)
+
+        # array
+        result = pd.Series(op(data, other))
+        expected = pd.Series(op(data._data, other), dtype="boolean")
+        # propagate NAs
+        expected[data._mask] = pd.NA
+
+        tm.assert_series_equal(result, expected)
+
+        # series
+        s = pd.Series(data)
+        result = op(s, other)
+
+        expected = pd.Series(data._data)
+        expected = op(expected, other)
+        expected = expected.astype("boolean")
+        # propagate NAs
+        expected[data._mask] = pd.NA
+
+        tm.assert_series_equal(result, expected)
+
+    def test_compare_scalar(self, data, all_compare_operators):
+        op_name = all_compare_operators
+        self._compare_other(data, op_name, True)
+
+    def test_compare_array(self, data, all_compare_operators):
+        op_name = all_compare_operators
+        other = pd.array([True] * len(data), dtype="boolean")
+        self._compare_other(data, op_name, other)
+        other = np.array([True] * len(data))
+        self._compare_other(data, op_name, other)
+        other = pd.Series([True] * len(data))
+        self._compare_other(data, op_name, other)
+
+    @pytest.mark.parametrize("other", [True, False, pd.NA])
+    def test_scalar(self, other, all_compare_operators):
+        op = self.get_op_from_name(all_compare_operators)
+        a = pd.array([True, False, None], dtype="boolean")
+
+        result = op(a, other)
+
+        if other is pd.NA:
+            expected = pd.array([None, None, None], dtype="boolean")
+        else:
+            values = op(a._data, other)
+            expected = BooleanArray(values, a._mask, copy=True)
+        tm.assert_extension_array_equal(result, expected)
+
+        # ensure we haven't mutated anything inplace
+        result[0] = None
+        tm.assert_extension_array_equal(
+            a, pd.array([True, False, None], dtype="boolean")
+        )
+
+    def test_array(self, all_compare_operators):
+        op = self.get_op_from_name(all_compare_operators)
+        a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
+        b = pd.array([True, False, None] * 3, dtype="boolean")
+
+        result = op(a, b)
+
+        values = op(a._data, b._data)
+        mask = a._mask | b._mask
+        expected = BooleanArray(values, mask)
+        tm.assert_extension_array_equal(result, expected)
+
+        # ensure we haven't mutated anything inplace
+        result[0] = None
+        tm.assert_extension_array_equal(
+            a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
+        )
+        tm.assert_extension_array_equal(
+            b, pd.array([True, False, None] * 3, dtype="boolean")
+        )