From 6a56a49a735a6f987836c0456086bdaad0ad42c9 Mon Sep 17 00:00:00 2001 From: Brock Date: Wed, 7 Oct 2020 16:59:29 -0700 Subject: [PATCH] REF/TYP: use OpsMixin for logical methods --- pandas/core/arraylike.py | 31 ++++++++++++++++++++++++++ pandas/core/ops/__init__.py | 27 ++-------------------- pandas/core/ops/methods.py | 25 ++++++++++++++------- pandas/core/series.py | 10 +++++++++ pandas/tests/extension/arrow/arrays.py | 13 ++++------- 5 files changed, 64 insertions(+), 42 deletions(-) diff --git a/pandas/core/arraylike.py b/pandas/core/arraylike.py index 1fba022f2a1de..185e9197e01fe 100644 --- a/pandas/core/arraylike.py +++ b/pandas/core/arraylike.py @@ -8,6 +8,7 @@ from pandas.errors import AbstractMethodError +from pandas.core.ops import roperator from pandas.core.ops.common import unpack_zerodim_and_defer @@ -41,3 +42,33 @@ def __gt__(self, other): @unpack_zerodim_and_defer("__ge__") def __ge__(self, other): return self._cmp_method(other, operator.ge) + + # ------------------------------------------------------------- + # Logical Methods + + def _logical_method(self, other, op): + raise AbstractMethodError(self) + + @unpack_zerodim_and_defer("__and__") + def __and__(self, other): + return self._logical_method(other, operator.and_) + + @unpack_zerodim_and_defer("__rand__") + def __rand__(self, other): + return self._logical_method(other, roperator.rand_) + + @unpack_zerodim_and_defer("__or__") + def __or__(self, other): + return self._logical_method(other, operator.or_) + + @unpack_zerodim_and_defer("__ror__") + def __ror__(self, other): + return self._logical_method(other, roperator.ror_) + + @unpack_zerodim_and_defer("__xor__") + def __xor__(self, other): + return self._logical_method(other, operator.xor) + + @unpack_zerodim_and_defer("__rxor__") + def __rxor__(self, other): + return self._logical_method(other, roperator.rxor) diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index 84319b69d9a35..ae21f13ea3f49 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -280,7 +280,7 @@ def dispatch_to_series(left, right, func, axis: Optional[int] = None): # Series -def _align_method_SERIES(left: "Series", right, align_asobject: bool = False): +def align_method_SERIES(left: "Series", right, align_asobject: bool = False): """ align lhs and rhs Series """ # ToDo: Different from align_method_FRAME, list, tuple and ndarray # are not coerced here @@ -311,7 +311,7 @@ def arith_method_SERIES(cls, op, special): @unpack_zerodim_and_defer(op_name) def wrapper(left, right): res_name = get_op_result_name(left, right) - left, right = _align_method_SERIES(left, right) + left, right = align_method_SERIES(left, right) lvalues = extract_array(left, extract_numpy=True) rvalues = extract_array(right, extract_numpy=True) @@ -323,29 +323,6 @@ def wrapper(left, right): return wrapper -def bool_method_SERIES(cls, op, special): - """ - Wrapper function for Series arithmetic operations, to avoid - code duplication. - """ - assert special # non-special uses flex_method_SERIES - op_name = _get_op_name(op, special) - - @unpack_zerodim_and_defer(op_name) - def wrapper(self, other): - res_name = get_op_result_name(self, other) - self, other = _align_method_SERIES(self, other, align_asobject=True) - - lvalues = extract_array(self, extract_numpy=True) - rvalues = extract_array(other, extract_numpy=True) - - res_values = logical_op(lvalues, rvalues, op) - return self._construct_result(res_values, name=res_name) - - wrapper.__name__ = op_name - return wrapper - - def flex_method_SERIES(cls, op, special): assert not special # "special" also means "not flex" name = _get_op_name(op, special) diff --git a/pandas/core/ops/methods.py b/pandas/core/ops/methods.py index 2b117d5e22186..70fd814423c7f 100644 --- a/pandas/core/ops/methods.py +++ b/pandas/core/ops/methods.py @@ -46,7 +46,6 @@ def _get_method_wrappers(cls): from pandas.core.ops import ( arith_method_FRAME, arith_method_SERIES, - bool_method_SERIES, comp_method_FRAME, flex_comp_method_FRAME, flex_method_SERIES, @@ -58,7 +57,7 @@ def _get_method_wrappers(cls): comp_flex = flex_method_SERIES arith_special = arith_method_SERIES comp_special = None - bool_special = bool_method_SERIES + bool_special = None elif issubclass(cls, ABCDataFrame): arith_flex = arith_method_FRAME comp_flex = flex_comp_method_FRAME @@ -118,13 +117,23 @@ def f(self, other): ) ) - new_methods.update( - dict( - __iand__=_wrap_inplace_method(new_methods["__and__"]), - __ior__=_wrap_inplace_method(new_methods["__or__"]), - __ixor__=_wrap_inplace_method(new_methods["__xor__"]), + if bool_method is None: + # Series gets bool_method via OpsMixin + new_methods.update( + dict( + __iand__=_wrap_inplace_method(cls.__and__), + __ior__=_wrap_inplace_method(cls.__or__), + __ixor__=_wrap_inplace_method(cls.__xor__), + ) + ) + else: + new_methods.update( + dict( + __iand__=_wrap_inplace_method(new_methods["__and__"]), + __ior__=_wrap_inplace_method(new_methods["__or__"]), + __ixor__=_wrap_inplace_method(new_methods["__xor__"]), + ) ) - ) _add_methods(cls, new_methods=new_methods) diff --git a/pandas/core/series.py b/pandas/core/series.py index 5cc163807fac6..0852f1b650ae6 100644 --- a/pandas/core/series.py +++ b/pandas/core/series.py @@ -4978,6 +4978,16 @@ def _cmp_method(self, other, op): return self._construct_result(res_values, name=res_name) + def _logical_method(self, other, op): + res_name = ops.get_op_result_name(self, other) + self, other = ops.align_method_SERIES(self, other, align_asobject=True) + + lvalues = extract_array(self, extract_numpy=True) + rvalues = extract_array(other, extract_numpy=True) + + res_values = ops.logical_op(lvalues, rvalues, op) + return self._construct_result(res_values, name=res_name) + Series._add_numeric_operations() diff --git a/pandas/tests/extension/arrow/arrays.py b/pandas/tests/extension/arrow/arrays.py index 5e930b7b22f30..04ce705690cf3 100644 --- a/pandas/tests/extension/arrow/arrays.py +++ b/pandas/tests/extension/arrow/arrays.py @@ -21,6 +21,7 @@ register_extension_dtype, take, ) +from pandas.core.arraylike import OpsMixin @register_extension_dtype @@ -67,7 +68,7 @@ def construct_array_type(cls) -> Type["ArrowStringArray"]: return ArrowStringArray -class ArrowExtensionArray(ExtensionArray): +class ArrowExtensionArray(OpsMixin, ExtensionArray): _data: pa.ChunkedArray @classmethod @@ -109,7 +110,7 @@ def astype(self, dtype, copy=True): def dtype(self): return self._dtype - def _boolean_op(self, other, op): + def _logical_method(self, other, op): if not isinstance(other, type(self)): raise NotImplementedError() @@ -122,13 +123,7 @@ def __eq__(self, other): if not isinstance(other, type(self)): return False - return self._boolean_op(other, operator.eq) - - def __and__(self, other): - return self._boolean_op(other, operator.and_) - - def __or__(self, other): - return self._boolean_op(other, operator.or_) + return self._logical_method(other, operator.eq) @property def nbytes(self) -> int: