Skip to content

Commit 60a7d2a

Browse files
authored
REF/TYP: use OpsMixin for logical methods (#36964)
1 parent d759308 commit 60a7d2a

File tree

5 files changed

+64
-42
lines changed

5 files changed

+64
-42
lines changed

pandas/core/arraylike.py

+31
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
from pandas.errors import AbstractMethodError
1010

11+
from pandas.core.ops import roperator
1112
from pandas.core.ops.common import unpack_zerodim_and_defer
1213

1314

@@ -41,3 +42,33 @@ def __gt__(self, other):
4142
@unpack_zerodim_and_defer("__ge__")
4243
def __ge__(self, other):
4344
return self._cmp_method(other, operator.ge)
45+
46+
# -------------------------------------------------------------
47+
# Logical Methods
48+
49+
def _logical_method(self, other, op):
50+
raise AbstractMethodError(self)
51+
52+
@unpack_zerodim_and_defer("__and__")
53+
def __and__(self, other):
54+
return self._logical_method(other, operator.and_)
55+
56+
@unpack_zerodim_and_defer("__rand__")
57+
def __rand__(self, other):
58+
return self._logical_method(other, roperator.rand_)
59+
60+
@unpack_zerodim_and_defer("__or__")
61+
def __or__(self, other):
62+
return self._logical_method(other, operator.or_)
63+
64+
@unpack_zerodim_and_defer("__ror__")
65+
def __ror__(self, other):
66+
return self._logical_method(other, roperator.ror_)
67+
68+
@unpack_zerodim_and_defer("__xor__")
69+
def __xor__(self, other):
70+
return self._logical_method(other, operator.xor)
71+
72+
@unpack_zerodim_and_defer("__rxor__")
73+
def __rxor__(self, other):
74+
return self._logical_method(other, roperator.rxor)

pandas/core/ops/__init__.py

+2-25
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,7 @@ def dispatch_to_series(left, right, func, axis: Optional[int] = None):
280280
# Series
281281

282282

283-
def _align_method_SERIES(left: "Series", right, align_asobject: bool = False):
283+
def align_method_SERIES(left: "Series", right, align_asobject: bool = False):
284284
""" align lhs and rhs Series """
285285
# ToDo: Different from align_method_FRAME, list, tuple and ndarray
286286
# are not coerced here
@@ -311,7 +311,7 @@ def arith_method_SERIES(cls, op, special):
311311
@unpack_zerodim_and_defer(op_name)
312312
def wrapper(left, right):
313313
res_name = get_op_result_name(left, right)
314-
left, right = _align_method_SERIES(left, right)
314+
left, right = align_method_SERIES(left, right)
315315

316316
lvalues = extract_array(left, extract_numpy=True)
317317
rvalues = extract_array(right, extract_numpy=True)
@@ -323,29 +323,6 @@ def wrapper(left, right):
323323
return wrapper
324324

325325

326-
def bool_method_SERIES(cls, op, special):
327-
"""
328-
Wrapper function for Series arithmetic operations, to avoid
329-
code duplication.
330-
"""
331-
assert special # non-special uses flex_method_SERIES
332-
op_name = _get_op_name(op, special)
333-
334-
@unpack_zerodim_and_defer(op_name)
335-
def wrapper(self, other):
336-
res_name = get_op_result_name(self, other)
337-
self, other = _align_method_SERIES(self, other, align_asobject=True)
338-
339-
lvalues = extract_array(self, extract_numpy=True)
340-
rvalues = extract_array(other, extract_numpy=True)
341-
342-
res_values = logical_op(lvalues, rvalues, op)
343-
return self._construct_result(res_values, name=res_name)
344-
345-
wrapper.__name__ = op_name
346-
return wrapper
347-
348-
349326
def flex_method_SERIES(cls, op, special):
350327
assert not special # "special" also means "not flex"
351328
name = _get_op_name(op, special)

pandas/core/ops/methods.py

+17-8
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,6 @@ def _get_method_wrappers(cls):
4646
from pandas.core.ops import (
4747
arith_method_FRAME,
4848
arith_method_SERIES,
49-
bool_method_SERIES,
5049
comp_method_FRAME,
5150
flex_comp_method_FRAME,
5251
flex_method_SERIES,
@@ -58,7 +57,7 @@ def _get_method_wrappers(cls):
5857
comp_flex = flex_method_SERIES
5958
arith_special = arith_method_SERIES
6059
comp_special = None
61-
bool_special = bool_method_SERIES
60+
bool_special = None
6261
elif issubclass(cls, ABCDataFrame):
6362
arith_flex = arith_method_FRAME
6463
comp_flex = flex_comp_method_FRAME
@@ -118,13 +117,23 @@ def f(self, other):
118117
)
119118
)
120119

121-
new_methods.update(
122-
dict(
123-
__iand__=_wrap_inplace_method(new_methods["__and__"]),
124-
__ior__=_wrap_inplace_method(new_methods["__or__"]),
125-
__ixor__=_wrap_inplace_method(new_methods["__xor__"]),
120+
if bool_method is None:
121+
# Series gets bool_method via OpsMixin
122+
new_methods.update(
123+
dict(
124+
__iand__=_wrap_inplace_method(cls.__and__),
125+
__ior__=_wrap_inplace_method(cls.__or__),
126+
__ixor__=_wrap_inplace_method(cls.__xor__),
127+
)
128+
)
129+
else:
130+
new_methods.update(
131+
dict(
132+
__iand__=_wrap_inplace_method(new_methods["__and__"]),
133+
__ior__=_wrap_inplace_method(new_methods["__or__"]),
134+
__ixor__=_wrap_inplace_method(new_methods["__xor__"]),
135+
)
126136
)
127-
)
128137

129138
_add_methods(cls, new_methods=new_methods)
130139

pandas/core/series.py

+10
Original file line numberDiff line numberDiff line change
@@ -4978,6 +4978,16 @@ def _cmp_method(self, other, op):
49784978

49794979
return self._construct_result(res_values, name=res_name)
49804980

4981+
def _logical_method(self, other, op):
4982+
res_name = ops.get_op_result_name(self, other)
4983+
self, other = ops.align_method_SERIES(self, other, align_asobject=True)
4984+
4985+
lvalues = extract_array(self, extract_numpy=True)
4986+
rvalues = extract_array(other, extract_numpy=True)
4987+
4988+
res_values = ops.logical_op(lvalues, rvalues, op)
4989+
return self._construct_result(res_values, name=res_name)
4990+
49814991

49824992
Series._add_numeric_operations()
49834993

pandas/tests/extension/arrow/arrays.py

+4-9
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
register_extension_dtype,
2222
take,
2323
)
24+
from pandas.core.arraylike import OpsMixin
2425

2526

2627
@register_extension_dtype
@@ -67,7 +68,7 @@ def construct_array_type(cls) -> Type["ArrowStringArray"]:
6768
return ArrowStringArray
6869

6970

70-
class ArrowExtensionArray(ExtensionArray):
71+
class ArrowExtensionArray(OpsMixin, ExtensionArray):
7172
_data: pa.ChunkedArray
7273

7374
@classmethod
@@ -109,7 +110,7 @@ def astype(self, dtype, copy=True):
109110
def dtype(self):
110111
return self._dtype
111112

112-
def _boolean_op(self, other, op):
113+
def _logical_method(self, other, op):
113114
if not isinstance(other, type(self)):
114115
raise NotImplementedError()
115116

@@ -122,13 +123,7 @@ def __eq__(self, other):
122123
if not isinstance(other, type(self)):
123124
return False
124125

125-
return self._boolean_op(other, operator.eq)
126-
127-
def __and__(self, other):
128-
return self._boolean_op(other, operator.and_)
129-
130-
def __or__(self, other):
131-
return self._boolean_op(other, operator.or_)
126+
return self._logical_method(other, operator.eq)
132127

133128
@property
134129
def nbytes(self) -> int:

0 commit comments

Comments
 (0)