REF: use OpsMixin in EAs (pandas-dev#37049)

Author: jbrockmendel

pandas/core/arrays/boolean.py

@@ -6,7 +6,6 @@
 
 from pandas._libs import lib, missing as libmissing
 from pandas._typing import ArrayLike
-from pandas.compat import set_function_name
 from pandas.compat.numpy import function as nv
 
 from pandas.core.dtypes.common import (
@@ -23,7 +22,6 @@
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops
-from pandas.core.arraylike import OpsMixin
 
 from .masked import BaseMaskedArray, BaseMaskedDtype
 
@@ -203,7 +201,7 @@ def coerce_to_array(
     return values, mask
 
 
-class BooleanArray(OpsMixin, BaseMaskedArray):
+class BooleanArray(BaseMaskedArray):
     """
     Array of boolean (True/False) data with missing values.
 
@@ -561,48 +559,40 @@ def all(self, skipna: bool = True, **kwargs):
             else:
                 return self.dtype.na_value
 
-    @classmethod
-    def _create_logical_method(cls, op):
-        @ops.unpack_zerodim_and_defer(op.__name__)
-        def logical_method(self, other):
-
-            assert op.__name__ in {"or_", "ror_", "and_", "rand_", "xor", "rxor"}
-            other_is_booleanarray = isinstance(other, BooleanArray)
-            other_is_scalar = lib.is_scalar(other)
-            mask = None
-
-            if other_is_booleanarray:
-                other, mask = other._data, other._mask
-            elif is_list_like(other):
-                other = np.asarray(other, dtype="bool")
-                if other.ndim > 1:
-                    raise NotImplementedError(
-                        "can only perform ops with 1-d structures"
-                    )
-                other, mask = coerce_to_array(other, copy=False)
-            elif isinstance(other, np.bool_):
-                other = other.item()
-
-            if other_is_scalar and not (other is libmissing.NA or lib.is_bool(other)):
-                raise TypeError(
-                    "'other' should be pandas.NA or a bool. "
-                    f"Got {type(other).__name__} instead."
-                )
-
-            if not other_is_scalar and len(self) != len(other):
-                raise ValueError("Lengths must match to compare")
+    def _logical_method(self, other, op):
+
+        assert op.__name__ in {"or_", "ror_", "and_", "rand_", "xor", "rxor"}
+        other_is_booleanarray = isinstance(other, BooleanArray)
+        other_is_scalar = lib.is_scalar(other)
+        mask = None
+
+        if other_is_booleanarray:
+            other, mask = other._data, other._mask
+        elif is_list_like(other):
+            other = np.asarray(other, dtype="bool")
+            if other.ndim > 1:
+                raise NotImplementedError("can only perform ops with 1-d structures")
+            other, mask = coerce_to_array(other, copy=False)
+        elif isinstance(other, np.bool_):
+            other = other.item()
 
-            if op.__name__ in {"or_", "ror_"}:
-                result, mask = ops.kleene_or(self._data, other, self._mask, mask)
-            elif op.__name__ in {"and_", "rand_"}:
-                result, mask = ops.kleene_and(self._data, other, self._mask, mask)
-            elif op.__name__ in {"xor", "rxor"}:
-                result, mask = ops.kleene_xor(self._data, other, self._mask, mask)
+        if other_is_scalar and not (other is libmissing.NA or lib.is_bool(other)):
+            raise TypeError(
+                "'other' should be pandas.NA or a bool. "
+                f"Got {type(other).__name__} instead."
+            )
 
-            return BooleanArray(result, mask)
+        if not other_is_scalar and len(self) != len(other):
+            raise ValueError("Lengths must match to compare")
 
-        name = f"__{op.__name__}__"
-        return set_function_name(logical_method, name, cls)
+        if op.__name__ in {"or_", "ror_"}:
+            result, mask = ops.kleene_or(self._data, other, self._mask, mask)
+        elif op.__name__ in {"and_", "rand_"}:
+            result, mask = ops.kleene_and(self._data, other, self._mask, mask)
+        elif op.__name__ in {"xor", "rxor"}:
+            result, mask = ops.kleene_xor(self._data, other, self._mask, mask)
+
+        return BooleanArray(result, mask)
 
     def _cmp_method(self, other, op):
         from pandas.arrays import FloatingArray, IntegerArray
@@ -643,6 +633,50 @@ def _cmp_method(self, other, op):
 
         return BooleanArray(result, mask, copy=False)
 
+    def _arith_method(self, other, op):
+        mask = None
+        op_name = op.__name__
+
+        if isinstance(other, BooleanArray):
+            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")
+
+        # nans propagate
+        if mask is None:
+            mask = self._mask
+            if other is libmissing.NA:
+                mask |= True
+        else:
+            mask = self._mask | mask
+
+        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"
+            else:
+                dtype = "bool"
+            result = np.zeros(len(self._data), dtype=dtype)
+        else:
+            with np.errstate(all="ignore"):
+                result = op(self._data, other)
+
+        # divmod returns a tuple
+        if op_name == "divmod":
+            div, mod = result
+            return (
+                self._maybe_mask_result(div, mask, other, "floordiv"),
+                self._maybe_mask_result(mod, mask, other, "mod"),
+            )
+
+        return self._maybe_mask_result(result, mask, other, op_name)
+
     def _reduce(self, name: str, skipna: bool = True, **kwargs):
 
         if name in {"any", "all"}:
@@ -678,60 +712,3 @@ def _maybe_mask_result(self, result, mask, other, op_name: str):
         else:
             result[mask] = np.nan
             return result
-
-    @classmethod
-    def _create_arithmetic_method(cls, op):
-        op_name = op.__name__
-
-        @ops.unpack_zerodim_and_defer(op_name)
-        def boolean_arithmetic_method(self, other):
-            mask = None
-
-            if isinstance(other, BooleanArray):
-                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")
-
-            # nans propagate
-            if mask is None:
-                mask = self._mask
-                if other is libmissing.NA:
-                    mask |= True
-            else:
-                mask = self._mask | mask
-
-            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"
-                else:
-                    dtype = "bool"
-                result = np.zeros(len(self._data), dtype=dtype)
-            else:
-                with np.errstate(all="ignore"):
-                    result = op(self._data, other)
-
-            # divmod returns a tuple
-            if op_name == "divmod":
-                div, mod = result
-                return (
-                    self._maybe_mask_result(div, mask, other, "floordiv"),
-                    self._maybe_mask_result(mod, mask, other, "mod"),
-                )
-
-            return self._maybe_mask_result(result, mask, other, op_name)
-
-        name = f"__{op_name}__"
-        return set_function_name(boolean_arithmetic_method, name, cls)
-
-
-BooleanArray._add_logical_ops()
-BooleanArray._add_arithmetic_ops()

pandas/core/arrays/floating.py

@@ -6,7 +6,6 @@
 
 from pandas._libs import lib, missing as libmissing
 from pandas._typing import ArrayLike, DtypeObj
-from pandas.compat import set_function_name
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import cache_readonly
 
@@ -26,9 +25,7 @@
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops
-from pandas.core.arraylike import OpsMixin
 from pandas.core.ops import invalid_comparison
-from pandas.core.ops.common import unpack_zerodim_and_defer
 from pandas.core.tools.numeric import to_numeric
 
 from .masked import BaseMaskedArray, BaseMaskedDtype
@@ -202,7 +199,7 @@ def coerce_to_array(
     return values, mask
 
 
-class FloatingArray(OpsMixin, BaseMaskedArray):
+class FloatingArray(BaseMaskedArray):
     """
     Array of floating (optional missing) values.
 
@@ -479,83 +476,70 @@ def _maybe_mask_result(self, result, mask, other, op_name: str):
 
         return type(self)(result, mask, copy=False)
 
-    @classmethod
-    def _create_arithmetic_method(cls, op):
-        op_name = op.__name__
-
-        @unpack_zerodim_and_defer(op.__name__)
-        def floating_arithmetic_method(self, other):
-            from pandas.arrays import IntegerArray
-
-            omask = None
+    def _arith_method(self, other, op):
+        from pandas.arrays import IntegerArray
 
-            if getattr(other, "ndim", 0) > 1:
-                raise NotImplementedError("can only perform ops with 1-d structures")
+        omask = None
 
-            if isinstance(other, (IntegerArray, FloatingArray)):
-                other, omask = other._data, other._mask
+        if getattr(other, "ndim", 0) > 1:
+            raise NotImplementedError("can only perform ops with 1-d structures")
 
-            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")
-                if not (is_float_dtype(other) or is_integer_dtype(other)):
-                    raise TypeError("can only perform ops with numeric values")
+        if isinstance(other, (IntegerArray, FloatingArray)):
+            other, omask = other._data, other._mask
 
-            else:
-                if not (is_float(other) or is_integer(other) or other is libmissing.NA):
-                    raise TypeError("can only perform ops with numeric values")
+        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")
+            if not (is_float_dtype(other) or is_integer_dtype(other)):
+                raise TypeError("can only perform ops with numeric values")
 
-            if omask is None:
-                mask = self._mask.copy()
-                if other is libmissing.NA:
-                    mask |= True
-            else:
-                mask = self._mask | omask
-
-            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)
+        else:
+            if not (is_float(other) or is_integer(other) or other is libmissing.NA):
+                raise TypeError("can only perform ops with numeric values")
 
+        if omask is None:
+            mask = self._mask.copy()
             if other is libmissing.NA:
-                result = np.ones_like(self._data)
-            else:
-                with np.errstate(all="ignore"):
-                    result = op(self._data, other)
-
-            # divmod returns a tuple
-            if op_name == "divmod":
-                div, mod = result
-                return (
-                    self._maybe_mask_result(div, mask, other, "floordiv"),
-                    self._maybe_mask_result(mod, mask, other, "mod"),
-                )
-
-            return self._maybe_mask_result(result, mask, other, op_name)
-
-        name = f"__{op.__name__}__"
-        return set_function_name(floating_arithmetic_method, name, cls)
+                mask |= True
+        else:
+            mask = self._mask | omask
+
+        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)
 
+        if other is libmissing.NA:
+            result = np.ones_like(self._data)
+        else:
+            with np.errstate(all="ignore"):
+                result = op(self._data, other)
+
+        # divmod returns a tuple
+        if op.__name__ == "divmod":
+            div, mod = result
+            return (
+                self._maybe_mask_result(div, mask, other, "floordiv"),
+                self._maybe_mask_result(mod, mask, other, "mod"),
+            )
 
-FloatingArray._add_arithmetic_ops()
+        return self._maybe_mask_result(result, mask, other, op.__name__)
 
 
 _dtype_docstring = """