diff --git a/pandas/core/arrays/boolean.py b/pandas/core/arrays/boolean.py index dd750bce7842e..4dd117e407961 100644 --- a/pandas/core/arrays/boolean.py +++ b/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() diff --git a/pandas/core/arrays/floating.py b/pandas/core/arrays/floating.py index bbb5467d42d53..aa272f13b045c 100644 --- a/pandas/core/arrays/floating.py +++ b/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 = """ diff --git a/pandas/core/arrays/integer.py b/pandas/core/arrays/integer.py index 258a946536c2b..856b4bcbda048 100644 --- a/pandas/core/arrays/integer.py +++ b/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 = """ diff --git a/pandas/core/arrays/numpy_.py b/pandas/core/arrays/numpy_.py index c56cccf2e4a93..b5103fb7f9d5d 100644 --- a/pandas/core/arrays/numpy_.py +++ b/pandas/core/arrays/numpy_.py @@ -14,6 +14,7 @@ from pandas import compat from pandas.core import nanops, ops from pandas.core.array_algos import masked_reductions +from pandas.core.arraylike import OpsMixin from pandas.core.arrays._mixins import NDArrayBackedExtensionArray from pandas.core.arrays.base import ExtensionOpsMixin from pandas.core.strings.object_array import ObjectStringArrayMixin @@ -115,6 +116,7 @@ def itemsize(self) -> int: class PandasArray( + OpsMixin, NDArrayBackedExtensionArray, ExtensionOpsMixin, NDArrayOperatorsMixin, @@ -370,31 +372,32 @@ def to_numpy( def __invert__(self): return type(self)(~self._ndarray) - @classmethod - def _create_arithmetic_method(cls, op): + def _cmp_method(self, other, op): + if isinstance(other, PandasArray): + other = other._ndarray pd_op = ops.get_array_op(op) + result = pd_op(self._ndarray, other) - @ops.unpack_zerodim_and_defer(op.__name__) - def arithmetic_method(self, other): - if isinstance(other, cls): - other = other._ndarray - - result = pd_op(self._ndarray, other) + if op is divmod or op is ops.rdivmod: + a, b = result + if isinstance(a, np.ndarray): + # for e.g. op vs TimedeltaArray, we may already + # have an ExtensionArray, in which case we do not wrap + return self._wrap_ndarray_result(a), self._wrap_ndarray_result(b) + return a, b - if op is divmod or op is ops.rdivmod: - a, b = result - if isinstance(a, np.ndarray): - # for e.g. op vs TimedeltaArray, we may already - # have an ExtensionArray, in which case we do not wrap - return self._wrap_ndarray_result(a), self._wrap_ndarray_result(b) - return a, b + if isinstance(result, np.ndarray): + # for e.g. multiplication vs TimedeltaArray, we may already + # have an ExtensionArray, in which case we do not wrap + return self._wrap_ndarray_result(result) + return result - if isinstance(result, np.ndarray): - # for e.g. multiplication vs TimedeltaArray, we may already - # have an ExtensionArray, in which case we do not wrap - return self._wrap_ndarray_result(result) - return result + @classmethod + def _create_arithmetic_method(cls, op): + @ops.unpack_zerodim_and_defer(op.__name__) + def arithmetic_method(self, other): + return self._cmp_method(other, op) return compat.set_function_name(arithmetic_method, f"__{op.__name__}__", cls) @@ -415,4 +418,3 @@ def _wrap_ndarray_result(self, result: np.ndarray): PandasArray._add_arithmetic_ops() -PandasArray._add_comparison_ops() diff --git a/pandas/core/arrays/string_.py b/pandas/core/arrays/string_.py index 9ea34d4680748..553ba25270943 100644 --- a/pandas/core/arrays/string_.py +++ b/pandas/core/arrays/string_.py @@ -314,43 +314,46 @@ def memory_usage(self, deep: bool = False) -> int: return result + lib.memory_usage_of_objects(self._ndarray) return result - # Override parent because we have different return types. - @classmethod - def _create_arithmetic_method(cls, op): - # Note: this handles both arithmetic and comparison methods. + def _cmp_method(self, other, op): + from pandas.arrays import BooleanArray - @ops.unpack_zerodim_and_defer(op.__name__) - def method(self, other): - from pandas.arrays import BooleanArray + if isinstance(other, StringArray): + other = other._ndarray + + mask = isna(self) | isna(other) + valid = ~mask - assert op.__name__ in ops.ARITHMETIC_BINOPS | ops.COMPARISON_BINOPS + if not lib.is_scalar(other): + if len(other) != len(self): + # prevent improper broadcasting when other is 2D + raise ValueError( + f"Lengths of operands do not match: {len(self)} != {len(other)}" + ) - if isinstance(other, cls): - other = other._ndarray + other = np.asarray(other) + other = other[valid] - mask = isna(self) | isna(other) - valid = ~mask + if op.__name__ in ops.ARITHMETIC_BINOPS: + result = np.empty_like(self._ndarray, dtype="object") + result[mask] = StringDtype.na_value + result[valid] = op(self._ndarray[valid], other) + return StringArray(result) + else: + # logical + result = np.zeros(len(self._ndarray), dtype="bool") + result[valid] = op(self._ndarray[valid], other) + return BooleanArray(result, mask) - if not lib.is_scalar(other): - if len(other) != len(self): - # prevent improper broadcasting when other is 2D - raise ValueError( - f"Lengths of operands do not match: {len(self)} != {len(other)}" - ) + # Override parent because we have different return types. + @classmethod + def _create_arithmetic_method(cls, op): + # Note: this handles both arithmetic and comparison methods. - other = np.asarray(other) - other = other[valid] + assert op.__name__ in ops.ARITHMETIC_BINOPS | ops.COMPARISON_BINOPS - if op.__name__ in ops.ARITHMETIC_BINOPS: - result = np.empty_like(self._ndarray, dtype="object") - result[mask] = StringDtype.na_value - result[valid] = op(self._ndarray[valid], other) - return StringArray(result) - else: - # logical - result = np.zeros(len(self._ndarray), dtype="bool") - result[valid] = op(self._ndarray[valid], other) - return BooleanArray(result, mask) + @ops.unpack_zerodim_and_defer(op.__name__) + def method(self, other): + return self._cmp_method(other, op) return compat.set_function_name(method, f"__{op.__name__}__", cls) @@ -362,7 +365,6 @@ def _add_arithmetic_ops(cls): cls.__mul__ = cls._create_arithmetic_method(operator.mul) cls.__rmul__ = cls._create_arithmetic_method(ops.rmul) - _create_comparison_method = _create_arithmetic_method # ------------------------------------------------------------------------ # String methods interface _str_na_value = StringDtype.na_value @@ -418,4 +420,3 @@ def _str_map(self, f, na_value=None, dtype=None): StringArray._add_arithmetic_ops() -StringArray._add_comparison_ops()