REF: share NumericArray._arith_method with BooleanArray (#45849)

Author: jbrockmendel

pandas/core/arrays/boolean.py

@@ -382,42 +382,5 @@ def _logical_method(self, other, op):
         # expected "ndarray"
         return BooleanArray(result, mask)  # type: ignore[arg-type]
 
-    def _arith_method(self, other, op):
-        mask = None
-        op_name = op.__name__
-
-        if isinstance(other, BaseMaskedArray):
-            other, mask = other._data, other._mask
-
-        elif is_list_like(other):
-            other = np.asarray(other)
-            if other.ndim > 1:
-                raise NotImplementedError("can only perform ops with 1-d structures")
-            if len(self) != len(other):
-                raise ValueError("Lengths must match")
-
-        mask = self._propagate_mask(mask, other)
-
-        if other is libmissing.NA:
-            # if other is NA, the result will be all NA and we can't run the
-            # actual op, so we need to choose the resulting dtype manually
-            if op_name in {"floordiv", "rfloordiv", "mod", "rmod", "pow", "rpow"}:
-                dtype = "int8"
-            elif op_name in {"truediv", "rtruediv"}:
-                dtype = "float64"
-            else:
-                dtype = "bool"
-            result = np.zeros(len(self._data), dtype=dtype)
-        else:
-            if op_name in {"pow", "rpow"} and isinstance(other, np.bool_):
-                # Avoid DeprecationWarning: In future, it will be an error
-                #  for 'np.bool_' scalars to be interpreted as an index
-                other = bool(other)
-
-            with np.errstate(all="ignore"):
-                result = op(self._data, other)
-
-        return self._maybe_mask_result(result, mask, other, op_name)
-
     def __abs__(self):
         return self.copy()

pandas/core/arrays/masked.py

@@ -76,6 +76,7 @@
 from pandas.core.array_algos.quantile import quantile_with_mask
 from pandas.core.arraylike import OpsMixin
 from pandas.core.arrays import ExtensionArray
+from pandas.core.construction import ensure_wrapped_if_datetimelike
 from pandas.core.indexers import check_array_indexer
 from pandas.core.ops import invalid_comparison
 
@@ -593,6 +594,85 @@ def _propagate_mask(
             mask = self._mask | mask
         return mask
 
+    def _arith_method(self, other, op):
+        op_name = op.__name__
+        omask = None
+
+        if isinstance(other, BaseMaskedArray):
+            other, omask = other._data, other._mask
+
+        elif is_list_like(other):
+            if not isinstance(other, ExtensionArray):
+                other = np.asarray(other)
+            if other.ndim > 1:
+                raise NotImplementedError("can only perform ops with 1-d structures")
+
+        # We wrap the non-masked arithmetic logic used for numpy dtypes
+        #  in Series/Index arithmetic ops.
+        other = ops.maybe_prepare_scalar_for_op(other, (len(self),))
+        pd_op = ops.get_array_op(op)
+        other = ensure_wrapped_if_datetimelike(other)
+
+        if op_name in {"pow", "rpow"} and isinstance(other, np.bool_):
+            # Avoid DeprecationWarning: In future, it will be an error
+            #  for 'np.bool_' scalars to be interpreted as an index
+            #  e.g. test_array_scalar_like_equivalence
+            other = bool(other)
+
+        mask = self._propagate_mask(omask, other)
+
+        if other is libmissing.NA:
+            result = np.ones_like(self._data)
+            if self.dtype.kind == "b":
+                if op_name in {"floordiv", "rfloordiv", "mod", "rmod", "pow", "rpow"}:
+                    dtype = "int8"
+                elif op_name in {"truediv", "rtruediv"}:
+                    dtype = "float64"
+                else:
+                    dtype = "bool"
+                result = result.astype(dtype)
+            elif "truediv" in op_name and self.dtype.kind != "f":
+                # The actual data here doesn't matter since the mask
+                #  will be all-True, but since this is division, we want
+                #  to end up with floating dtype.
+                result = result.astype(np.float64)
+        else:
+            # Make sure we do this before the "pow" mask checks
+            #  to get an expected exception message on shape mismatch.
+            if self.dtype.kind in ["i", "u"] and op_name in ["floordiv", "mod"]:
+                # TODO(GH#30188) ATM we don't match the behavior of non-masked
+                #  types with respect to floordiv-by-zero
+                pd_op = op
+
+            elif self.dtype.kind == "b" and (
+                "div" in op_name or "pow" in op_name or "mod" in op_name
+            ):
+                # TODO(GH#41165): should these be disallowed?
+                pd_op = op
+
+            with np.errstate(all="ignore"):
+                result = pd_op(self._data, other)
+
+        if op_name == "pow":
+            # 1 ** x is 1.
+            mask = np.where((self._data == 1) & ~self._mask, False, mask)
+            # x ** 0 is 1.
+            if omask is not None:
+                mask = np.where((other == 0) & ~omask, False, mask)
+            elif other is not libmissing.NA:
+                mask = np.where(other == 0, False, mask)
+
+        elif op_name == "rpow":
+            # 1 ** x is 1.
+            if omask is not None:
+                mask = np.where((other == 1) & ~omask, False, mask)
+            elif other is not libmissing.NA:
+                mask = np.where(other == 1, False, mask)
+            # x ** 0 is 1.
+            mask = np.where((self._data == 0) & ~self._mask, False, mask)
+
+        return self._maybe_mask_result(result, mask, other, op_name)
+
     def _cmp_method(self, other, op) -> BooleanArray:
         from pandas.core.arrays import BooleanArray
 

pandas/core/arrays/numeric.py

@@ -23,19 +23,15 @@
     is_bool_dtype,
     is_float_dtype,
     is_integer_dtype,
-    is_list_like,
     is_object_dtype,
     is_string_dtype,
     pandas_dtype,
 )
 
-from pandas.core import ops
-from pandas.core.arrays.base import ExtensionArray
 from pandas.core.arrays.masked import (
     BaseMaskedArray,
     BaseMaskedDtype,
 )
-from pandas.core.construction import ensure_wrapped_if_datetimelike
 
 if TYPE_CHECKING:
     import pyarrow
@@ -214,65 +210,6 @@ def _from_sequence_of_strings(
         scalars = to_numeric(strings, errors="raise")
         return cls._from_sequence(scalars, dtype=dtype, copy=copy)
 
-    def _arith_method(self, other, op):
-        op_name = op.__name__
-        omask = None
-
-        if isinstance(other, BaseMaskedArray):
-            other, omask = other._data, other._mask
-
-        elif is_list_like(other):
-            if not isinstance(other, ExtensionArray):
-                other = np.asarray(other)
-            if other.ndim > 1:
-                raise NotImplementedError("can only perform ops with 1-d structures")
-
-        # We wrap the non-masked arithmetic logic used for numpy dtypes
-        #  in Series/Index arithmetic ops.
-        other = ops.maybe_prepare_scalar_for_op(other, (len(self),))
-        pd_op = ops.get_array_op(op)
-        other = ensure_wrapped_if_datetimelike(other)
-
-        mask = self._propagate_mask(omask, other)
-
-        if other is libmissing.NA:
-            result = np.ones_like(self._data)
-            if "truediv" in op_name and self.dtype.kind != "f":
-                # The actual data here doesn't matter since the mask
-                #  will be all-True, but since this is division, we want
-                #  to end up with floating dtype.
-                result = result.astype(np.float64)
-        else:
-            # Make sure we do this before the "pow" mask checks
-            #  to get an expected exception message on shape mismatch.
-            if self.dtype.kind in ["i", "u"] and op_name in ["floordiv", "mod"]:
-                # ATM we don't match the behavior of non-masked types with
-                #  respect to floordiv-by-zero
-                pd_op = op
-
-            with np.errstate(all="ignore"):
-                result = pd_op(self._data, other)
-
-        if op_name == "pow":
-            # 1 ** x is 1.
-            mask = np.where((self._data == 1) & ~self._mask, False, mask)
-            # x ** 0 is 1.
-            if omask is not None:
-                mask = np.where((other == 0) & ~omask, False, mask)
-            elif other is not libmissing.NA:
-                mask = np.where(other == 0, False, mask)
-
-        elif op_name == "rpow":
-            # 1 ** x is 1.
-            if omask is not None:
-                mask = np.where((other == 1) & ~omask, False, mask)
-            elif other is not libmissing.NA:
-                mask = np.where(other == 1, False, mask)
-            # x ** 0 is 1.
-            mask = np.where((self._data == 0) & ~self._mask, False, mask)
-
-        return self._maybe_mask_result(result, mask, other, op_name)
-
     _HANDLED_TYPES = (np.ndarray, numbers.Number)
 
     def __neg__(self):

pandas/tests/arrays/boolean/test_arithmetic.py

@@ -71,9 +71,7 @@ def test_div(left_array, right_array):
     [
         "floordiv",
         "mod",
-        pytest.param(
-            "pow", marks=pytest.mark.xfail(reason="TODO follow int8 behaviour? GH34686")
-        ),
+        "pow",
     ],
 )
 def test_op_int8(left_array, right_array, opname):

pandas/tests/arrays/masked/test_arithmetic.py

@@ -153,22 +153,29 @@ def test_error_len_mismatch(data, all_arithmetic_operators):
 
     other = [scalar] * (len(data) - 1)
 
+    err = ValueError
     msg = "|".join(
         [
             r"operands could not be broadcast together with shapes \(3,\) \(4,\)",
             r"operands could not be broadcast together with shapes \(4,\) \(3,\)",
         ]
     )
-
-    if data.dtype.kind == "b":
-        msg = "Lengths must match"
+    if data.dtype.kind == "b" and all_arithmetic_operators.strip("_") in [
+        "sub",
+        "rsub",
+    ]:
+        err = TypeError
+        msg = (
+            r"numpy boolean subtract, the `\-` operator, is not supported, use "
+            r"the bitwise_xor, the `\^` operator, or the logical_xor function instead"
+        )
 
     for other in [other, np.array(other)]:
-        with pytest.raises(ValueError, match=msg):
+        with pytest.raises(err, match=msg):
             op(data, other)
 
         s = pd.Series(data)
-        with pytest.raises(ValueError, match=msg):
+        with pytest.raises(err, match=msg):
             op(s, other)