TYP/REF: use _cmp_method in EAs (#36954)

Author: jbrockmendel

pandas/core/arrays/boolean.py

@@ -23,6 +23,7 @@
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops
+from pandas.core.arraylike import OpsMixin
 
 from .masked import BaseMaskedArray, BaseMaskedDtype
 
@@ -202,7 +203,7 @@ def coerce_to_array(
     return values, mask
 
 
-class BooleanArray(BaseMaskedArray):
+class BooleanArray(OpsMixin, BaseMaskedArray):
     """
     Array of boolean (True/False) data with missing values.
 
@@ -603,52 +604,44 @@ def logical_method(self, other):
         name = f"__{op.__name__}__"
         return set_function_name(logical_method, name, cls)
 
-    @classmethod
-    def _create_comparison_method(cls, op):
-        @ops.unpack_zerodim_and_defer(op.__name__)
-        def cmp_method(self, other):
-            from pandas.arrays import FloatingArray, IntegerArray
+    def _cmp_method(self, other, op):
+        from pandas.arrays import FloatingArray, IntegerArray
 
-            if isinstance(other, (IntegerArray, FloatingArray)):
-                return NotImplemented
+        if isinstance(other, (IntegerArray, FloatingArray)):
+            return NotImplemented
 
-            mask = None
+        mask = None
 
-            if isinstance(other, BooleanArray):
-                other, mask = other._data, other._mask
+        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 to compare")
+        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 to compare")
 
-            if other is libmissing.NA:
-                # numpy does not handle pd.NA well as "other" scalar (it returns
-                # a scalar False instead of an array)
-                result = np.zeros_like(self._data)
-                mask = np.ones_like(self._data)
-            else:
-                # numpy will show a DeprecationWarning on invalid elementwise
-                # comparisons, this will raise in the future
-                with warnings.catch_warnings():
-                    warnings.filterwarnings("ignore", "elementwise", FutureWarning)
-                    with np.errstate(all="ignore"):
-                        result = op(self._data, other)
-
-                # nans propagate
-                if mask is None:
-                    mask = self._mask.copy()
-                else:
-                    mask = self._mask | mask
+        if other is libmissing.NA:
+            # numpy does not handle pd.NA well as "other" scalar (it returns
+            # a scalar False instead of an array)
+            result = np.zeros_like(self._data)
+            mask = np.ones_like(self._data)
+        else:
+            # numpy will show a DeprecationWarning on invalid elementwise
+            # comparisons, this will raise in the future
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", "elementwise", FutureWarning)
+                with np.errstate(all="ignore"):
+                    result = op(self._data, other)
 
-            return BooleanArray(result, mask, copy=False)
+            # nans propagate
+            if mask is None:
+                mask = self._mask.copy()
+            else:
+                mask = self._mask | mask
 
-        name = f"__{op.__name__}"
-        return set_function_name(cmp_method, name, cls)
+        return BooleanArray(result, mask, copy=False)
 
     def _reduce(self, name: str, skipna: bool = True, **kwargs):
 
@@ -741,5 +734,4 @@ def boolean_arithmetic_method(self, other):
 
 
 BooleanArray._add_logical_ops()
-BooleanArray._add_comparison_ops()
 BooleanArray._add_arithmetic_ops()

pandas/core/arrays/floating.py

@@ -26,6 +26,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
@@ -201,7 +202,7 @@ def coerce_to_array(
     return values, mask
 
 
-class FloatingArray(BaseMaskedArray):
+class FloatingArray(OpsMixin, BaseMaskedArray):
     """
     Array of floating (optional missing) values.
 
@@ -398,58 +399,48 @@ def astype(self, dtype, copy: bool = True) -> ArrayLike:
     def _values_for_argsort(self) -> np.ndarray:
         return self._data
 
-    @classmethod
-    def _create_comparison_method(cls, op):
-        op_name = op.__name__
+    def _cmp_method(self, other, op):
+        from pandas.arrays import BooleanArray, IntegerArray
 
-        @unpack_zerodim_and_defer(op.__name__)
-        def cmp_method(self, other):
-            from pandas.arrays import BooleanArray, IntegerArray
+        mask = None
 
-            mask = None
+        if isinstance(other, (BooleanArray, IntegerArray, FloatingArray)):
+            other, mask = other._data, other._mask
 
-            if isinstance(other, (BooleanArray, IntegerArray, FloatingArray)):
-                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")
 
-            elif is_list_like(other):
-                other = np.asarray(other)
-                if other.ndim > 1:
-                    raise NotImplementedError(
-                        "can only perform ops with 1-d structures"
-                    )
+        if other is libmissing.NA:
+            # numpy does not handle pd.NA well as "other" scalar (it returns
+            # a scalar False instead of an array)
+            # This may be fixed by NA.__array_ufunc__. Revisit this check
+            # once that's implemented.
+            result = np.zeros(self._data.shape, dtype="bool")
+            mask = np.ones(self._data.shape, dtype="bool")
+        else:
+            with warnings.catch_warnings():
+                # numpy may show a FutureWarning:
+                #     elementwise comparison failed; returning scalar instead,
+                #     but in the future will perform elementwise comparison
+                # before returning NotImplemented. We fall back to the correct
+                # behavior today, so that should be fine to ignore.
+                warnings.filterwarnings("ignore", "elementwise", FutureWarning)
+                with np.errstate(all="ignore"):
+                    method = getattr(self._data, f"__{op.__name__}__")
+                    result = method(other)
 
-            if other is libmissing.NA:
-                # numpy does not handle pd.NA well as "other" scalar (it returns
-                # a scalar False instead of an array)
-                # This may be fixed by NA.__array_ufunc__. Revisit this check
-                # once that's implemented.
-                result = np.zeros(self._data.shape, dtype="bool")
-                mask = np.ones(self._data.shape, dtype="bool")
-            else:
-                with warnings.catch_warnings():
-                    # numpy may show a FutureWarning:
-                    #     elementwise comparison failed; returning scalar instead,
-                    #     but in the future will perform elementwise comparison
-                    # before returning NotImplemented. We fall back to the correct
-                    # behavior today, so that should be fine to ignore.
-                    warnings.filterwarnings("ignore", "elementwise", FutureWarning)
-                    with np.errstate(all="ignore"):
-                        method = getattr(self._data, f"__{op_name}__")
-                        result = method(other)
-
-                    if result is NotImplemented:
-                        result = invalid_comparison(self._data, other, op)
-
-            # nans propagate
-            if mask is None:
-                mask = self._mask.copy()
-            else:
-                mask = self._mask | mask
+                if result is NotImplemented:
+                    result = invalid_comparison(self._data, other, op)
 
-            return BooleanArray(result, mask)
+        # nans propagate
+        if mask is None:
+            mask = self._mask.copy()
+        else:
+            mask = self._mask | mask
 
-        name = f"__{op.__name__}__"
-        return set_function_name(cmp_method, name, cls)
+        return BooleanArray(result, mask)
 
     def sum(self, skipna=True, min_count=0, **kwargs):
         nv.validate_sum((), kwargs)
@@ -565,7 +556,6 @@ def floating_arithmetic_method(self, other):
 
 
 FloatingArray._add_arithmetic_ops()
-FloatingArray._add_comparison_ops()
 
 
 _dtype_docstring = """

pandas/core/arrays/integer.py

@@ -26,6 +26,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
@@ -265,7 +266,7 @@ def coerce_to_array(
     return values, mask
 
 
-class IntegerArray(BaseMaskedArray):
+class IntegerArray(OpsMixin, BaseMaskedArray):
     """
     Array of integer (optional missing) values.
 
@@ -493,60 +494,50 @@ def _values_for_argsort(self) -> np.ndarray:
             data[self._mask] = data.min() - 1
         return data
 
-    @classmethod
-    def _create_comparison_method(cls, op):
-        op_name = op.__name__
+    def _cmp_method(self, other, op):
+        from pandas.core.arrays import BaseMaskedArray, BooleanArray
 
-        @unpack_zerodim_and_defer(op.__name__)
-        def cmp_method(self, other):
-            from pandas.core.arrays import BaseMaskedArray, BooleanArray
+        mask = None
 
-            mask = None
+        if isinstance(other, BaseMaskedArray):
+            other, mask = other._data, other._mask
 
-            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 to compare")
+        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 to compare")
+
+        if other is libmissing.NA:
+            # numpy does not handle pd.NA well as "other" scalar (it returns
+            # a scalar False instead of an array)
+            # This may be fixed by NA.__array_ufunc__. Revisit this check
+            # once that's implemented.
+            result = np.zeros(self._data.shape, dtype="bool")
+            mask = np.ones(self._data.shape, dtype="bool")
+        else:
+            with warnings.catch_warnings():
+                # numpy may show a FutureWarning:
+                #     elementwise comparison failed; returning scalar instead,
+                #     but in the future will perform elementwise comparison
+                # before returning NotImplemented. We fall back to the correct
+                # behavior today, so that should be fine to ignore.
+                warnings.filterwarnings("ignore", "elementwise", FutureWarning)
+                with np.errstate(all="ignore"):
+                    method = getattr(self._data, f"__{op.__name__}__")
+                    result = method(other)
 
-            if other is libmissing.NA:
-                # numpy does not handle pd.NA well as "other" scalar (it returns
-                # a scalar False instead of an array)
-                # This may be fixed by NA.__array_ufunc__. Revisit this check
-                # once that's implemented.
-                result = np.zeros(self._data.shape, dtype="bool")
-                mask = np.ones(self._data.shape, dtype="bool")
-            else:
-                with warnings.catch_warnings():
-                    # numpy may show a FutureWarning:
-                    #     elementwise comparison failed; returning scalar instead,
-                    #     but in the future will perform elementwise comparison
-                    # before returning NotImplemented. We fall back to the correct
-                    # behavior today, so that should be fine to ignore.
-                    warnings.filterwarnings("ignore", "elementwise", FutureWarning)
-                    with np.errstate(all="ignore"):
-                        method = getattr(self._data, f"__{op_name}__")
-                        result = method(other)
-
-                    if result is NotImplemented:
-                        result = invalid_comparison(self._data, other, op)
-
-            # nans propagate
-            if mask is None:
-                mask = self._mask.copy()
-            else:
-                mask = self._mask | mask
+                if result is NotImplemented:
+                    result = invalid_comparison(self._data, other, op)
 
-            return BooleanArray(result, mask)
+        # nans propagate
+        if mask is None:
+            mask = self._mask.copy()
+        else:
+            mask = self._mask | mask
 
-        name = f"__{op.__name__}__"
-        return set_function_name(cmp_method, name, cls)
+        return BooleanArray(result, mask)
 
     def sum(self, skipna=True, min_count=0, **kwargs):
         nv.validate_sum((), kwargs)
@@ -669,7 +660,6 @@ def integer_arithmetic_method(self, other):
 
 
 IntegerArray._add_arithmetic_ops()
-IntegerArray._add_comparison_ops()
 
 
 _dtype_docstring = """