def arithmetic_op( left: Union[np.ndarray, ABCExtensionArray], right: Any, op, str_rep: str, eval_kwargs: Mapping[str, bool], ): """ 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. str_rep : str 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 keep_null_freq = isinstance( right, ( ABCDatetimeIndex, ABCDatetimeArray, ABCTimedeltaIndex, ABCTimedeltaArray, Timestamp, ), ) # 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) 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: with np.errstate(all="ignore"): res_values = na_arithmetic_op(lvalues, rvalues, op, str_rep, eval_kwargs) return res_values
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 """ # 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(lvalues.shape, dtype=bool) else: res_values = np.zeros(lvalues.shape, dtype=bool) elif is_object_dtype(lvalues.dtype): res_values = comp_method_OBJECT_ARRAY(op, lvalues, rvalues) else: op_name = f"__{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): typ = type(rvalues) raise TypeError(f"Could not compare {typ} type with Series") return res_values
def comparison_op( left: ArrayLike, right: Any, op, str_rep: Optional[str] = None, ) -> ArrayLike: """ 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 ------- ndarray or ExtensionArray """ # 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(lvalues.shape, dtype=bool) else: res_values = np.zeros(lvalues.shape, dtype=bool) elif is_object_dtype(lvalues.dtype): res_values = comp_method_OBJECT_ARRAY(op, lvalues, rvalues) else: with np.errstate(all="ignore"): res_values = na_arithmetic_op(lvalues, rvalues, op, str_rep, is_cmp=True) return res_values
def arithmetic_op(left: Union[np.ndarray, ABCExtensionArray], right: Any, op, str_rep: str): """ 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. str_rep : str 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 # NB: We assume that extract_array has already been called # on `left` and `right`. lvalues = left rvalues = right rvalues = maybe_upcast_for_op(rvalues, lvalues.shape) if should_extension_dispatch(left, rvalues) or isinstance( rvalues, (ABCTimedeltaArray, ABCDatetimeArray, Timestamp, Timedelta)): # TimedeltaArray, DatetimeArray, and Timestamp are included here # because they have `freq` attribute which is handled correctly # by dispatch_to_extension_op. # Timedelta is included because numexpr will fail on it, see GH#31457 res_values = dispatch_to_extension_op(op, lvalues, rvalues) else: with np.errstate(all="ignore"): res_values = na_arithmetic_op(lvalues, rvalues, op, str_rep) return res_values
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 """ 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_logical_op(lvalues, rvalues, op) res_values = filler(res_values) # type: ignore return res_values