From 9a617e3eada0a664ba9d5a3d27b21bf28d2659c7 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Thu, 12 Sep 2019 15:25:33 -0700 Subject: [PATCH 01/10] REF: implement logical and comparison array ops --- pandas/core/ops/__init__.py | 124 ++-------------------------- pandas/core/ops/array_ops.py | 155 ++++++++++++++++++++++++++++++++++- 2 files changed, 157 insertions(+), 122 deletions(-) diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index 016feff7e3beb..731b48163f0a1 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -9,14 +9,11 @@ import numpy as np -from pandas._libs import Timedelta, Timestamp, lib, ops as libops +from pandas._libs import Timedelta, Timestamp, lib from pandas.errors import NullFrequencyError from pandas.util._decorators import Appender -from pandas.core.dtypes.cast import construct_1d_object_array_from_listlike from pandas.core.dtypes.common import ( - ensure_object, - is_bool_dtype, is_datetime64_dtype, is_extension_array_dtype, is_integer_dtype, @@ -40,14 +37,15 @@ from pandas._typing import ArrayLike from pandas.core.construction import array, extract_array -from pandas.core.ops.array_ops import comp_method_OBJECT_ARRAY, define_na_arithmetic_op +from pandas.core.ops.array_ops import comp_method_OBJECT_ARRAY # noqa:F401 +from pandas.core.ops.array_ops import comparison_op, define_na_arithmetic_op, logical_op from pandas.core.ops.docstrings import ( _arith_doc_FRAME, _flex_comp_doc_FRAME, _make_flex_doc, _op_descriptions, ) -from pandas.core.ops.invalid import invalid_comparison +from pandas.core.ops.invalid import invalid_comparison # noqa:F401 from pandas.core.ops.methods import ( # noqa:F401 add_flex_arithmetic_methods, add_special_arithmetic_methods, @@ -695,46 +693,10 @@ def wrapper(self, other): if isinstance(other, ABCSeries) and not self._indexed_same(other): raise ValueError("Can only compare identically-labeled Series objects") - other = lib.item_from_zerodim(other) - if isinstance(other, list): - # TODO: same for tuples? - other = np.asarray(other) - - if isinstance(other, (np.ndarray, ABCExtensionArray, ABCIndexClass)): - # TODO: make this treatment consistent across ops and classes. - # We are not catching all listlikes here (e.g. frozenset, tuple) - # The ambiguous case is object-dtype. See GH#27803 - if len(self) != len(other): - raise ValueError("Lengths must match to compare") - lvalues = extract_array(self, extract_numpy=True) rvalues = extract_array(other, extract_numpy=True) - if should_extension_dispatch(lvalues, rvalues): - res_values = dispatch_to_extension_op(op, lvalues, rvalues) - - elif is_scalar(rvalues) and isna(rvalues): - # numpy does not like comparisons vs None - if op is operator.ne: - res_values = np.ones(len(lvalues), dtype=bool) - else: - res_values = np.zeros(len(lvalues), dtype=bool) - - elif is_object_dtype(lvalues.dtype): - res_values = comp_method_OBJECT_ARRAY(op, lvalues, rvalues) - - else: - op_name = "__{op}__".format(op=op.__name__) - method = getattr(lvalues, op_name) - with np.errstate(all="ignore"): - res_values = method(rvalues) - - if res_values is NotImplemented: - res_values = invalid_comparison(lvalues, rvalues, op) - if is_scalar(res_values): - raise TypeError( - "Could not compare {typ} type with Series".format(typ=type(rvalues)) - ) + res_values = comparison_op(lvalues, rvalues, op) result = self._constructor(res_values, index=self.index) result = finalizer(result) @@ -755,58 +717,7 @@ def _bool_method_SERIES(cls, op, special): """ op_name = _get_op_name(op, special) - def na_op(x, y): - try: - result = op(x, y) - except TypeError: - assert not isinstance(y, (list, ABCSeries, ABCIndexClass)) - if isinstance(y, np.ndarray): - # bool-bool dtype operations should be OK, should not get here - assert not (is_bool_dtype(x.dtype) and is_bool_dtype(y.dtype)) - x = ensure_object(x) - y = ensure_object(y) - result = libops.vec_binop(x, y, op) - else: - # let null fall thru - assert lib.is_scalar(y) - if not isna(y): - y = bool(y) - try: - result = libops.scalar_binop(x, y, op) - except ( - TypeError, - ValueError, - AttributeError, - OverflowError, - NotImplementedError, - ): - raise TypeError( - "cannot compare a dtyped [{dtype}] array " - "with a scalar of type [{typ}]".format( - dtype=x.dtype, typ=type(y).__name__ - ) - ) - - return result - - fill_int = lambda x: x - - def fill_bool(x, left=None): - # if `left` is specifically not-boolean, we do not cast to bool - if x.dtype.kind in ["c", "f", "O"]: - # dtypes that can hold NA - mask = isna(x) - if mask.any(): - x = x.astype(object) - x[mask] = False - - if left is None or is_bool_dtype(left.dtype): - x = x.astype(bool) - return x - def wrapper(self, other): - is_self_int_dtype = is_integer_dtype(self.dtype) - self, other = _align_method_SERIES(self, other, align_asobject=True) res_name = get_op_result_name(self, other) @@ -822,33 +733,10 @@ def wrapper(self, other): # Defer to DataFrame implementation; fail early return NotImplemented - other = lib.item_from_zerodim(other) - if is_list_like(other) and not hasattr(other, "dtype"): - # e.g. list, tuple - other = construct_1d_object_array_from_listlike(other) - lvalues = extract_array(self, extract_numpy=True) rvalues = extract_array(other, extract_numpy=True) - if should_extension_dispatch(self, rvalues): - res_values = dispatch_to_extension_op(op, lvalues, rvalues) - - else: - if isinstance(rvalues, (ABCSeries, ABCIndexClass, np.ndarray)): - is_other_int_dtype = is_integer_dtype(rvalues.dtype) - rvalues = rvalues if is_other_int_dtype else fill_bool(rvalues, lvalues) - - else: - # i.e. scalar - is_other_int_dtype = lib.is_integer(rvalues) - - # For int vs int `^`, `|`, `&` are bitwise operators and return - # integer dtypes. Otherwise these are boolean ops - filler = fill_int if is_self_int_dtype and is_other_int_dtype else fill_bool - - res_values = na_op(lvalues, rvalues) - res_values = filler(res_values) - + res_values = logical_op(lvalues, rvalues, op) result = self._constructor(res_values, index=self.index, name=res_name) return finalizer(result) diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index f5f6d77676f1f..a2f67e6accbae 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -2,20 +2,35 @@ Functions for arithmetic and comparison operations on NumPy arrays and ExtensionArrays. """ +import operator + import numpy as np -from pandas._libs import ops as libops +from pandas._libs import lib, ops as libops from pandas.core.dtypes.cast import ( construct_1d_object_array_from_listlike, find_common_type, maybe_upcast_putmask, ) -from pandas.core.dtypes.common import is_object_dtype, is_scalar -from pandas.core.dtypes.generic import ABCIndex, ABCSeries -from pandas.core.dtypes.missing import notna +from pandas.core.dtypes.common import ( + ensure_object, + is_bool_dtype, + is_integer_dtype, + is_list_like, + is_object_dtype, + is_scalar, +) +from pandas.core.dtypes.generic import ( + ABCExtensionArray, + ABCIndex, + ABCIndexClass, + ABCSeries, +) +from pandas.core.dtypes.missing import isna, notna from pandas.core.ops import missing +from pandas.core.ops.invalid import invalid_comparison from pandas.core.ops.roperator import rpow @@ -126,3 +141,135 @@ def na_op(x, y): return missing.dispatch_fill_zeros(op, x, y, result) return na_op + + +def comparison_op(left, right, op): + from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op + + # NB: We assume extract_array has already been called on left and right + lvalues = left + rvalues = right + + rvalues = lib.item_from_zerodim(rvalues) + if isinstance(rvalues, list): + # TODO: same for tuples? + rvalues = np.asarray(rvalues) + + if isinstance(rvalues, (np.ndarray, ABCExtensionArray, ABCIndexClass)): + # TODO: make this treatment consistent across ops and classes. + # We are not catching all listlikes here (e.g. frozenset, tuple) + # The ambiguous case is object-dtype. See GH#27803 + if len(lvalues) != len(rvalues): + raise ValueError("Lengths must match to compare") + + if should_extension_dispatch(lvalues, rvalues): + res_values = dispatch_to_extension_op(op, lvalues, rvalues) + + elif is_scalar(rvalues) and isna(rvalues): + # numpy does not like comparisons vs None + if op is operator.ne: + res_values = np.ones(len(lvalues), dtype=bool) + else: + res_values = np.zeros(len(lvalues), dtype=bool) + + elif is_object_dtype(lvalues.dtype): + res_values = comp_method_OBJECT_ARRAY(op, lvalues, rvalues) + + else: + op_name = "__{op}__".format(op=op.__name__) + method = getattr(lvalues, op_name) + with np.errstate(all="ignore"): + res_values = method(rvalues) + + if res_values is NotImplemented: + res_values = invalid_comparison(lvalues, rvalues, op) + if is_scalar(res_values): + raise TypeError( + "Could not compare {typ} type with Series".format(typ=type(rvalues)) + ) + + return res_values + + +def logical_op(left, right, op): + from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op + + def na_op(x, y): + try: + result = op(x, y) + except TypeError: + if isinstance(y, np.ndarray): + # bool-bool dtype operations should be OK, should not get here + assert not (is_bool_dtype(x.dtype) and is_bool_dtype(y.dtype)) + x = ensure_object(x) + y = ensure_object(y) + result = libops.vec_binop(x, y, op) + else: + # let null fall thru + assert lib.is_scalar(y) + if not isna(y): + y = bool(y) + try: + result = libops.scalar_binop(x, y, op) + except ( + TypeError, + ValueError, + AttributeError, + OverflowError, + NotImplementedError, + ): + raise TypeError( + "cannot compare a dtyped [{dtype}] array " + "with a scalar of type [{typ}]".format( + dtype=x.dtype, typ=type(y).__name__ + ) + ) + + return result + + fill_int = lambda x: x + + def fill_bool(x, left=None): + # if `left` is specifically not-boolean, we do not cast to bool + if x.dtype.kind in ["c", "f", "O"]: + # dtypes that can hold NA + mask = isna(x) + if mask.any(): + x = x.astype(object) + x[mask] = False + + if left is None or is_bool_dtype(left.dtype): + x = x.astype(bool) + return x + + is_self_int_dtype = is_integer_dtype(left.dtype) + + right = lib.item_from_zerodim(right) + if is_list_like(right) and not hasattr(right, "dtype"): + # e.g. list, tuple + right = construct_1d_object_array_from_listlike(right) + + # NB: We assume extract_array has already been called on left and right + lvalues = left + rvalues = right + + if should_extension_dispatch(lvalues, rvalues): + res_values = dispatch_to_extension_op(op, lvalues, rvalues) + + else: + if isinstance(rvalues, np.ndarray): + is_other_int_dtype = is_integer_dtype(rvalues.dtype) + rvalues = rvalues if is_other_int_dtype else fill_bool(rvalues, lvalues) + + else: + # i.e. scalar + is_other_int_dtype = lib.is_integer(rvalues) + + # For int vs int `^`, `|`, `&` are bitwise operators and return + # integer dtypes. Otherwise these are boolean ops + filler = fill_int if is_self_int_dtype and is_other_int_dtype else fill_bool + + res_values = na_op(lvalues, rvalues) + res_values = filler(res_values) + + return res_values From 56dff20313393d7ea1a4358b779b64fc4141393c Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Fri, 13 Sep 2019 09:10:29 -0700 Subject: [PATCH 02/10] implement arithmetic_op --- pandas/core/ops/__init__.py | 31 ++------------------------ pandas/core/ops/array_ops.py | 43 +++++++++++++++++++++++++++++++++++- 2 files changed, 44 insertions(+), 30 deletions(-) diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index 64893b198f552..f7b5e7d90927c 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -24,24 +24,20 @@ ) from pandas.core.dtypes.generic import ( ABCDataFrame, - ABCDatetimeArray, - ABCDatetimeIndex, ABCExtensionArray, ABCIndexClass, ABCSeries, ABCSparseSeries, - ABCTimedeltaArray, - ABCTimedeltaIndex, ) from pandas.core.dtypes.missing import isna, notna from pandas._typing import ArrayLike from pandas.core.construction import array, extract_array from pandas.core.ops.array_ops import ( + arithmetic_op, comparison_op, define_na_arithmetic_op, logical_op, - na_arithmetic_op, ) from pandas.core.ops.array_ops import comp_method_OBJECT_ARRAY # noqa:F401 from pandas.core.ops.docstrings import ( @@ -637,30 +633,7 @@ def wrapper(left, right): left, right = _align_method_SERIES(left, right) res_name = get_op_result_name(left, right) - keep_null_freq = isinstance( - right, - ( - ABCDatetimeIndex, - ABCDatetimeArray, - ABCTimedeltaIndex, - ABCTimedeltaArray, - Timestamp, - ), - ) - - lvalues = extract_array(left, extract_numpy=True) - rvalues = extract_array(right, extract_numpy=True) - - rvalues = maybe_upcast_for_op(rvalues, lvalues.shape) - - if should_extension_dispatch(left, rvalues) or isinstance( - rvalues, (ABCTimedeltaArray, ABCDatetimeArray, Timestamp) - ): - result = dispatch_to_extension_op(op, lvalues, rvalues, keep_null_freq) - - else: - with np.errstate(all="ignore"): - result = na_arithmetic_op(lvalues, rvalues, op, str_rep, eval_kwargs) + result = arithmetic_op(left, right, op, str_rep, eval_kwargs) # We do not pass dtype to ensure that the Series constructor # does inference in the case where `result` has object-dtype. diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index 5adbbc1d6b6bc..4abb50485a524 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -6,7 +6,7 @@ import numpy as np -from pandas._libs import lib, ops as libops +from pandas._libs import Timestamp, lib, ops as libops from pandas.core.dtypes.cast import ( construct_1d_object_array_from_listlike, @@ -22,13 +22,18 @@ is_scalar, ) from pandas.core.dtypes.generic import ( + ABCDatetimeArray, + ABCDatetimeIndex, ABCExtensionArray, ABCIndex, ABCIndexClass, ABCSeries, + ABCTimedeltaArray, + ABCTimedeltaIndex, ) from pandas.core.dtypes.missing import isna, notna +from pandas.core.construction import extract_array from pandas.core.ops import missing from pandas.core.ops.invalid import invalid_comparison from pandas.core.ops.roperator import rpow @@ -149,6 +154,42 @@ def na_arithmetic_op(left, right, op, str_rep, eval_kwargs): return missing.dispatch_fill_zeros(op, left, right, result) +def arithmetic_op(left, right, op, str_rep, eval_kwargs): + + from pandas.core.ops import ( + maybe_upcast_for_op, + should_extension_dispatch, + dispatch_to_extension_op, + ) + + keep_null_freq = isinstance( + right, + ( + ABCDatetimeIndex, + ABCDatetimeArray, + ABCTimedeltaIndex, + ABCTimedeltaArray, + Timestamp, + ), + ) + + lvalues = extract_array(left, extract_numpy=True) + rvalues = extract_array(right, extract_numpy=True) + + rvalues = maybe_upcast_for_op(rvalues, lvalues.shape) + + if should_extension_dispatch(left, rvalues) or isinstance( + rvalues, (ABCTimedeltaArray, ABCDatetimeArray, Timestamp) + ): + res_values = dispatch_to_extension_op(op, lvalues, rvalues, keep_null_freq) + + else: + with np.errstate(all="ignore"): + res_values = na_arithmetic_op(lvalues, rvalues, op, str_rep, eval_kwargs) + + return res_values + + def comparison_op(left, right, op): from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op From 148a8e85abd8b6423c7b7dbabfd59c1a16c5c78a Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Fri, 13 Sep 2019 15:17:00 -0700 Subject: [PATCH 03/10] add comments, types --- pandas/core/ops/__init__.py | 3 ++- pandas/core/ops/array_ops.py | 21 +++++++++++++++++---- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index f7b5e7d90927c..0edb59b487ab7 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -633,7 +633,8 @@ def wrapper(left, right): left, right = _align_method_SERIES(left, right) res_name = get_op_result_name(left, right) - result = arithmetic_op(left, right, op, str_rep, eval_kwargs) + lvalues = extract_array(left, extract_numpy=True) + result = arithmetic_op(lvalues, right, op, str_rep, eval_kwargs) # We do not pass dtype to ensure that the Series constructor # does inference in the case where `result` has object-dtype. diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index 4abb50485a524..dd7f0562cbc56 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -3,6 +3,7 @@ ExtensionArrays. """ import operator +from typing import Any, Optional, Union import numpy as np @@ -154,7 +155,13 @@ def na_arithmetic_op(left, right, op, str_rep, eval_kwargs): return missing.dispatch_fill_zeros(op, left, right, result) -def arithmetic_op(left, right, op, str_rep, eval_kwargs): +def arithmetic_op( + left: Union[ABCExtensionArray, np.ndarrray], + right: Any, + op, + str_rep: Optional[str], + eval_kwargs: dict, +): from pandas.core.ops import ( maybe_upcast_for_op, @@ -173,7 +180,10 @@ def arithmetic_op(left, right, op, str_rep, eval_kwargs): ), ) - lvalues = extract_array(left, extract_numpy=True) + # NB: We assume that extract_array has already been called on `left`, but + # cannot make the same assumption about `right`. This is because we need + # to define `keep_null_freq` before calling extract_array on it. + lvalues = left rvalues = extract_array(right, extract_numpy=True) rvalues = maybe_upcast_for_op(rvalues, lvalues.shape) @@ -181,6 +191,9 @@ def arithmetic_op(left, right, op, str_rep, eval_kwargs): if should_extension_dispatch(left, rvalues) or isinstance( rvalues, (ABCTimedeltaArray, ABCDatetimeArray, Timestamp) ): + # TimedeltaArray, DatetimeArray, and Timestamp are included here + # because they have `freq` attribute which is handled correctly + # by dispatch_to_extension_op. res_values = dispatch_to_extension_op(op, lvalues, rvalues, keep_null_freq) else: @@ -190,7 +203,7 @@ def arithmetic_op(left, right, op, str_rep, eval_kwargs): return res_values -def comparison_op(left, right, op): +def comparison_op(left: Union[ABCExtensionArray, np.ndarrray], right: Any, op): from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op # NB: We assume extract_array has already been called on left and right @@ -238,7 +251,7 @@ def comparison_op(left, right, op): return res_values -def logical_op(left, right, op): +def logical_op(left: Union[ABCExtensionArray, np.ndarrray], right: Any, op): from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op def na_op(x, y): From fcf973562760dcbf7761d57bea9f8a2c2b3968ae Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Fri, 13 Sep 2019 15:48:04 -0700 Subject: [PATCH 04/10] typo fixup --- pandas/core/ops/array_ops.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index dd7f0562cbc56..e333ee77c517e 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -156,7 +156,7 @@ def na_arithmetic_op(left, right, op, str_rep, eval_kwargs): def arithmetic_op( - left: Union[ABCExtensionArray, np.ndarrray], + left: Union[ABCExtensionArray, np.ndarray], right: Any, op, str_rep: Optional[str], @@ -203,7 +203,7 @@ def arithmetic_op( return res_values -def comparison_op(left: Union[ABCExtensionArray, np.ndarrray], right: Any, op): +def comparison_op(left: Union[ABCExtensionArray, np.ndarray], right: Any, op): from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op # NB: We assume extract_array has already been called on left and right @@ -251,7 +251,7 @@ def comparison_op(left: Union[ABCExtensionArray, np.ndarrray], right: Any, op): return res_values -def logical_op(left: Union[ABCExtensionArray, np.ndarrray], right: Any, op): +def logical_op(left: Union[ABCExtensionArray, np.ndarray], right: Any, op): from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op def na_op(x, y): From fec86dee6c076acc9425f0b4488f348a5c129426 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Fri, 13 Sep 2019 17:34:27 -0700 Subject: [PATCH 05/10] revert types --- pandas/core/ops/array_ops.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index e333ee77c517e..94c5a0fae480b 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -3,7 +3,6 @@ ExtensionArrays. """ import operator -from typing import Any, Optional, Union import numpy as np @@ -155,13 +154,7 @@ def na_arithmetic_op(left, right, op, str_rep, eval_kwargs): return missing.dispatch_fill_zeros(op, left, right, result) -def arithmetic_op( - left: Union[ABCExtensionArray, np.ndarray], - right: Any, - op, - str_rep: Optional[str], - eval_kwargs: dict, -): +def arithmetic_op(left, right, op, str_rep, eval_kwargs): from pandas.core.ops import ( maybe_upcast_for_op, @@ -203,7 +196,7 @@ def arithmetic_op( return res_values -def comparison_op(left: Union[ABCExtensionArray, np.ndarray], right: Any, op): +def comparison_op(left, right, op): from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op # NB: We assume extract_array has already been called on left and right @@ -251,7 +244,7 @@ def comparison_op(left: Union[ABCExtensionArray, np.ndarray], right: Any, op): return res_values -def logical_op(left: Union[ABCExtensionArray, np.ndarray], right: Any, op): +def logical_op(left, right, op): from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op def na_op(x, y): From 2abdccb828cf99075f23ebaf43c98bc03e3d5b27 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Tue, 17 Sep 2019 15:48:10 -0700 Subject: [PATCH 06/10] add types --- pandas/core/ops/array_ops.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index 94c5a0fae480b..b5807f18b8d54 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -3,6 +3,7 @@ ExtensionArrays. """ import operator +from typing import Any, Dict, Union import numpy as np @@ -154,7 +155,13 @@ def na_arithmetic_op(left, right, op, str_rep, eval_kwargs): return missing.dispatch_fill_zeros(op, left, right, result) -def arithmetic_op(left, right, op, str_rep, eval_kwargs): +def arithmetic_op( + left: Union[np.ndarray, ABCExtensionArray], + right: Any, + op, + str_rep: str, + eval_kwargs: Dict[str, str], +) -> Union[np.ndarray, ABCExtensionArray]: from pandas.core.ops import ( maybe_upcast_for_op, @@ -196,7 +203,9 @@ def arithmetic_op(left, right, op, str_rep, eval_kwargs): return res_values -def comparison_op(left, right, op): +def comparison_op( + left: Union[np.ndarray, ABCExtensionArray], right: Any, op +) -> Union[np.ndarray, ABCExtensionArray]: from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op # NB: We assume extract_array has already been called on left and right @@ -244,7 +253,9 @@ def comparison_op(left, right, op): return res_values -def logical_op(left, right, op): +def logical_op( + left: Union[np.ndarray, ABCExtensionArray], right: Any, op +) -> Union[np.ndarray, ABCExtensionArray]: from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op def na_op(x, y): From 121d7838d7e684d3e1d105c66ed2bf1224efc067 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Tue, 17 Sep 2019 15:54:21 -0700 Subject: [PATCH 07/10] docstrings --- pandas/core/ops/array_ops.py | 45 ++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index b5807f18b8d54..c62bcd888c4c7 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -162,6 +162,22 @@ def arithmetic_op( str_rep: str, eval_kwargs: Dict[str, str], ) -> Union[np.ndarray, ABCExtensionArray]: + """ + Evaluate an arithmetic operation `+`, `-`, `*`, `/`, `//`, `%`, `**`, ... + + Parameters + ---------- + left : np.ndarray or ExtensionArray + right : object + Cannot be a DataFrame or Index. Series is *not* excluded. + op : {operator.add, operator.sub, ...} + Or one of the reversed variants from roperator. + + Returns + ------- + ndarrray or ExtensionArray + Or a 2-tuple of these in the case of divmod or rdivmod. + """ from pandas.core.ops import ( maybe_upcast_for_op, @@ -206,6 +222,20 @@ def arithmetic_op( def comparison_op( left: Union[np.ndarray, ABCExtensionArray], right: Any, op ) -> Union[np.ndarray, ABCExtensionArray]: + """ + Evaluate a comparison operation `=`, `!=`, `>=`, `>`, `<=`, or `<`. + + Parameters + ---------- + left : np.ndarray or ExtensionArray + right : object + Cannot be a DataFrame, Series, or Index. + op : {operator.eq, operator.ne, operator.gt, operator.ge, operator.lt, operator.le} + + Returns + ------- + ndarrray or ExtensionArray + """ from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op # NB: We assume extract_array has already been called on left and right @@ -256,6 +286,21 @@ def comparison_op( def logical_op( left: Union[np.ndarray, ABCExtensionArray], right: Any, op ) -> Union[np.ndarray, ABCExtensionArray]: + """ + Evaluate a logical operation `|`, `&`, or `^`. + + Parameters + ---------- + left : np.ndarray or ExtensionArray + right : object + Cannot be a DataFrame, Series, or Index. + op : {operator.and_, operator.or_, operator.xor} + Or one of the reversed variants from roperator. + + Returns + ------- + ndarrray or ExtensionArray + """ from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op def na_op(x, y): From 267c7cacce57923a8b39a4cd85c0baaaa305bf1e Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Tue, 17 Sep 2019 16:10:57 -0700 Subject: [PATCH 08/10] ignore type --- pandas/core/ops/array_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index c62bcd888c4c7..3b15400f97c45 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -379,6 +379,6 @@ def fill_bool(x, left=None): filler = fill_int if is_self_int_dtype and is_other_int_dtype else fill_bool res_values = na_op(lvalues, rvalues) - res_values = filler(res_values) + res_values = filler(res_values) # type: ignore return res_values From 0b5aa34747dfc25a3629c94543fc0850334e8ea1 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Tue, 17 Sep 2019 16:11:42 -0700 Subject: [PATCH 09/10] revert technically-incorrect type --- pandas/core/ops/array_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index 3b15400f97c45..95da803192d63 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -161,7 +161,7 @@ def arithmetic_op( op, str_rep: str, eval_kwargs: Dict[str, str], -) -> Union[np.ndarray, ABCExtensionArray]: +): """ Evaluate an arithmetic operation `+`, `-`, `*`, `/`, `//`, `%`, `**`, ... From 8ced97bf8f827ea7a4d2e5a3d1d8c76ce7ee6bed Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Wed, 18 Sep 2019 07:19:59 -0700 Subject: [PATCH 10/10] REF: move na_op out --- pandas/core/ops/array_ops.py | 69 ++++++++++++++++++------------------ 1 file changed, 35 insertions(+), 34 deletions(-) diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index 95da803192d63..b72ef69ede199 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -283,6 +283,40 @@ def comparison_op( return res_values +def na_logical_op(x, y, op): + try: + result = op(x, y) + except TypeError: + if isinstance(y, np.ndarray): + # bool-bool dtype operations should be OK, should not get here + assert not (is_bool_dtype(x.dtype) and is_bool_dtype(y.dtype)) + x = ensure_object(x) + y = ensure_object(y) + result = libops.vec_binop(x, y, op) + else: + # let null fall thru + assert lib.is_scalar(y) + if not isna(y): + y = bool(y) + try: + result = libops.scalar_binop(x, y, op) + except ( + TypeError, + ValueError, + AttributeError, + OverflowError, + NotImplementedError, + ): + raise TypeError( + "cannot compare a dtyped [{dtype}] array " + "with a scalar of type [{typ}]".format( + dtype=x.dtype, typ=type(y).__name__ + ) + ) + + return result + + def logical_op( left: Union[np.ndarray, ABCExtensionArray], right: Any, op ) -> Union[np.ndarray, ABCExtensionArray]: @@ -303,39 +337,6 @@ def logical_op( """ from pandas.core.ops import should_extension_dispatch, dispatch_to_extension_op - def na_op(x, y): - try: - result = op(x, y) - except TypeError: - if isinstance(y, np.ndarray): - # bool-bool dtype operations should be OK, should not get here - assert not (is_bool_dtype(x.dtype) and is_bool_dtype(y.dtype)) - x = ensure_object(x) - y = ensure_object(y) - result = libops.vec_binop(x, y, op) - else: - # let null fall thru - assert lib.is_scalar(y) - if not isna(y): - y = bool(y) - try: - result = libops.scalar_binop(x, y, op) - except ( - TypeError, - ValueError, - AttributeError, - OverflowError, - NotImplementedError, - ): - raise TypeError( - "cannot compare a dtyped [{dtype}] array " - "with a scalar of type [{typ}]".format( - dtype=x.dtype, typ=type(y).__name__ - ) - ) - - return result - fill_int = lambda x: x def fill_bool(x, left=None): @@ -378,7 +379,7 @@ def fill_bool(x, left=None): # integer dtypes. Otherwise these are boolean ops filler = fill_int if is_self_int_dtype and is_other_int_dtype else fill_bool - res_values = na_op(lvalues, rvalues) + res_values = na_logical_op(lvalues, rvalues, op) res_values = filler(res_values) # type: ignore return res_values