def _make_compare(op): opname = '__{op}__'.format(op=op.__name__) def _evaluate_compare(self, other): # if we have a Categorical type, then must have the same # categories if isinstance(other, CategoricalIndex): other = other._values elif isinstance(other, Index): other = self._create_categorical( self, other._values, categories=self.categories, ordered=self.ordered) if isinstance(other, (ABCCategorical, np.ndarray, ABCSeries)): if len(self.values) != len(other): raise ValueError("Lengths must match to compare") if isinstance(other, ABCCategorical): if not self.values.is_dtype_equal(other): raise TypeError("categorical index comparisons must " "have the same categories and ordered " "attributes") result = op(self.values, other) if isinstance(result, ABCSeries): # Dispatch to pd.Categorical returned NotImplemented # and we got a Series back; down-cast to ndarray result = result.values return result return compat.set_function_name(_evaluate_compare, opname, cls)
def _make_comparison_op(op, cls): # TODO: share code with indexes.base version? Main difference is that # the block for MultiIndex was removed here. def cmp_method(self, other): if isinstance(other, ABCDataFrame): return NotImplemented if isinstance(other, (np.ndarray, ABCIndexClass, ABCSeries)): if other.ndim > 0 and len(self) != len(other): raise ValueError('Lengths must match to compare') if needs_i8_conversion(self) and needs_i8_conversion(other): # we may need to directly compare underlying # representations return self._evaluate_compare(other, op) # numpy will show a DeprecationWarning on invalid elementwise # comparisons, this will raise in the future with warnings.catch_warnings(record=True): with np.errstate(all='ignore'): result = op(self.values, np.asarray(other)) return result name = '__{name}__'.format(name=op.__name__) # TODO: docstring? return compat.set_function_name(cmp_method, name, cls)
def _create_comparison_method(cls, op): def cmp_method(self, other): op_name = op.__name__ mask = None if isinstance(other, IntegerArray): other, mask = other._data, other._mask elif is_list_like(other): other = np.asarray(other) if other.ndim > 0 and len(self) != len(other): raise ValueError('Lengths must match to compare') # numpy will show a DeprecationWarning on invalid elementwise # comparisons, this will raise in the future with warnings.catch_warnings(record=True): with np.errstate(all='ignore'): result = op(self._data, other) # nans propagate if mask is None: mask = self._mask else: mask = self._mask | mask result[mask] = True if op_name == 'ne' else False return result name = '__{name}__'.format(name=op.__name__) return set_function_name(cmp_method, name, cls)
def _make_compare(opname): def _evaluate_compare(self, other): # if we have a Categorical type, then must have the same # categories if isinstance(other, CategoricalIndex): other = other._values elif isinstance(other, Index): other = self._create_categorical( self, other._values, categories=self.categories, ordered=self.ordered) if isinstance(other, (ABCCategorical, np.ndarray, ABCSeries)): if len(self.values) != len(other): raise ValueError("Lengths must match to compare") if isinstance(other, ABCCategorical): if not self.values.is_dtype_equal(other): raise TypeError("categorical index comparisions must " "have the same categories and ordered " "attributes") return getattr(self.values, opname)(other) return compat.set_function_name(_evaluate_compare, opname, cls)
def _dt_array_cmp(cls, op): """ Wrap comparison operations to convert datetime-like to datetime64 """ opname = '__{name}__'.format(name=op.__name__) nat_result = True if opname == '__ne__' else False def wrapper(self, other): meth = getattr(dtl.DatetimeLikeArrayMixin, opname) if isinstance(other, (datetime, np.datetime64, compat.string_types)): if isinstance(other, (datetime, np.datetime64)): # GH#18435 strings get a pass from tzawareness compat self._assert_tzawareness_compat(other) try: other = _to_m8(other, tz=self.tz) except ValueError: # string that cannot be parsed to Timestamp return ops.invalid_comparison(self, other, op) result = meth(self, other) if isna(other): result.fill(nat_result) elif lib.is_scalar(other): return ops.invalid_comparison(self, other, op) else: if isinstance(other, list): # FIXME: This can break for object-dtype with mixed types other = type(self)(other) elif not isinstance(other, (np.ndarray, ABCIndexClass, ABCSeries)): # Following Timestamp convention, __eq__ is all-False # and __ne__ is all True, others raise TypeError. return ops.invalid_comparison(self, other, op) if is_object_dtype(other): result = op(self.astype('O'), np.array(other)) elif not (is_datetime64_dtype(other) or is_datetime64tz_dtype(other)): # e.g. is_timedelta64_dtype(other) return ops.invalid_comparison(self, other, op) else: self._assert_tzawareness_compat(other) result = meth(self, np.asarray(other)) result = com.values_from_object(result) # Make sure to pass an array to result[...]; indexing with # Series breaks with older version of numpy o_mask = np.array(isna(other)) if o_mask.any(): result[o_mask] = nat_result if self.hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _create_arithmetic_method(cls, op): def integer_arithmetic_method(self, other): op_name = op.__name__ mask = None if isinstance(other, (ABCSeries, ABCIndexClass)): # Rely on pandas to unbox and dispatch to us. return NotImplemented if getattr(other, 'ndim', 0) > 1: raise NotImplementedError( "can only perform ops with 1-d structures") if isinstance(other, IntegerArray): other, mask = other._data, other._mask elif getattr(other, 'ndim', None) == 0: other = other.item() elif is_list_like(other): other = np.asarray(other) if not other.ndim: other = other.item() elif other.ndim == 1: if not (is_float_dtype(other) or is_integer_dtype(other)): raise TypeError( "can only perform ops with numeric values") else: if not (is_float(other) or is_integer(other)): raise TypeError("can only perform ops with numeric values") # nans propagate if mask is None: mask = self._mask else: mask = self._mask | mask # 1 ** np.nan is 1. So we have to unmask those. if op_name == 'pow': mask = np.where(self == 1, False, mask) elif op_name == 'rpow': mask = np.where(other == 1, False, mask) with np.errstate(all='ignore'): result = op(self._data, other) # divmod returns a tuple if op_name == 'divmod': div, mod = result return (self._maybe_mask_result(div, mask, other, 'floordiv'), self._maybe_mask_result(mod, mask, other, 'mod')) return self._maybe_mask_result(result, mask, other, op_name) name = '__{name}__'.format(name=op.__name__) return set_function_name(integer_arithmetic_method, name, cls)
def _create_unary_method(cls, op): def sparse_unary_method(self): fill_value = op(np.array(self.fill_value)).item() values = op(self.sp_values) dtype = SparseDtype(values.dtype, fill_value) return cls._simple_new(values, self.sp_index, dtype) name = '__{name}__'.format(name=op.__name__) return compat.set_function_name(sparse_unary_method, name, cls)
def _create_method(cls, op, coerce_to_dtype=True): """ A class method that returns a method that will correspond to an operator for an ExtensionArray subclass, by dispatching to the relevant operator defined on the individual elements of the ExtensionArray. Parameters ---------- op : function An operator that takes arguments op(a, b) coerce_to_dtype : bool boolean indicating whether to attempt to convert the result to the underlying ExtensionArray dtype (default True) Returns ------- A method that can be bound to a method of a class Example ------- Given an ExtensionArray subclass called MyExtensionArray, use >>> __add__ = cls._create_method(operator.add) in the class definition of MyExtensionArray to create the operator for addition, that will be based on the operator implementation of the underlying elements of the ExtensionArray """ def _binop(self, other): def convert_values(param): if isinstance(param, ExtensionArray) or is_list_like(param): ovalues = param else: # Assume its an object ovalues = [param] * len(self) return ovalues lvalues = self rvalues = convert_values(other) # If the operator is not defined for the underlying objects, # a TypeError should be raised res = [op(a, b) for (a, b) in zip(lvalues, rvalues)] if coerce_to_dtype: try: res = self._from_sequence(res) except TypeError: pass return res op_name = ops._get_op_name(op, True) return set_function_name(_binop, op_name, cls)
def _td_index_cmp(opname, cls): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ def wrapper(self, other): result = getattr(TimedeltaArrayMixin, opname)(self, other) if is_bool_dtype(result): # support of bool dtype indexers return result return Index(result) return compat.set_function_name(wrapper, opname, cls)
def _period_array_cmp(cls, op): """ Wrap comparison operations to convert Period-like to PeriodDtype """ opname = '__{name}__'.format(name=op.__name__) nat_result = True if opname == '__ne__' else False def wrapper(self, other): op = getattr(self.asi8, opname) # We want to eventually defer to the Series or PeriodIndex (which will # return here with an unboxed PeriodArray). But before we do that, # we do a bit of validation on type (Period) and freq, so that our # error messages are sensible not_implemented = isinstance(other, (ABCSeries, ABCIndexClass)) if not_implemented: other = other._values if isinstance(other, Period): if other.freq != self.freq: msg = DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr) raise IncompatibleFrequency(msg) result = op(other.ordinal) elif isinstance(other, cls): if other.freq != self.freq: msg = DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr) raise IncompatibleFrequency(msg) if not_implemented: return NotImplemented result = op(other.asi8) mask = self._isnan | other._isnan if mask.any(): result[mask] = nat_result return result elif other is NaT: result = np.empty(len(self.asi8), dtype=bool) result.fill(nat_result) else: other = Period(other, freq=self.freq) result = op(other.ordinal) if self.hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _dt_array_cmp(opname, cls): """ Wrap comparison operations to convert datetime-like to datetime64 """ nat_result = True if opname == '__ne__' else False def wrapper(self, other): meth = getattr(dtl.DatetimeLikeArrayMixin, opname) if isinstance(other, (datetime, np.datetime64, compat.string_types)): if isinstance(other, datetime): # GH#18435 strings get a pass from tzawareness compat self._assert_tzawareness_compat(other) other = _to_m8(other, tz=self.tz) result = meth(self, other) if isna(other): result.fill(nat_result) else: if isinstance(other, list): other = type(self)(other) elif not isinstance(other, (np.ndarray, ABCIndexClass, ABCSeries)): # Following Timestamp convention, __eq__ is all-False # and __ne__ is all True, others raise TypeError. if opname == '__eq__': return np.zeros(shape=self.shape, dtype=bool) elif opname == '__ne__': return np.ones(shape=self.shape, dtype=bool) raise TypeError('%s type object %s' % (type(other), str(other))) if is_datetimelike(other): self._assert_tzawareness_compat(other) result = meth(self, np.asarray(other)) result = com._values_from_object(result) # Make sure to pass an array to result[...]; indexing with # Series breaks with older version of numpy o_mask = np.array(isna(other)) if o_mask.any(): result[o_mask] = nat_result if self.hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _td_array_cmp(cls, op): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ opname = '__{name}__'.format(name=op.__name__) nat_result = opname == '__ne__' def wrapper(self, other): other = lib.item_from_zerodim(other) if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)): return NotImplemented if _is_convertible_to_td(other) or other is NaT: try: other = Timedelta(other) except ValueError: # failed to parse as timedelta return ops.invalid_comparison(self, other, op) result = op(self.view('i8'), other.value) if isna(other): result.fill(nat_result) elif not is_list_like(other): return ops.invalid_comparison(self, other, op) elif len(other) != len(self): raise ValueError("Lengths must match") else: try: other = type(self)._from_sequence(other)._data except (ValueError, TypeError): return ops.invalid_comparison(self, other, op) result = op(self.view('i8'), other.view('i8')) result = com.values_from_object(result) o_mask = np.array(isna(other)) if o_mask.any(): result[o_mask] = nat_result if self._hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _td_array_cmp(cls, op): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ opname = '__{name}__'.format(name=op.__name__) nat_result = True if opname == '__ne__' else False meth = getattr(dtl.DatetimeLikeArrayMixin, opname) def wrapper(self, other): if _is_convertible_to_td(other) or other is NaT: try: other = _to_m8(other) except ValueError: # failed to parse as timedelta return ops.invalid_comparison(self, other, op) result = meth(self, other) if isna(other): result.fill(nat_result) elif not is_list_like(other): return ops.invalid_comparison(self, other, op) elif len(other) != len(self): raise ValueError("Lengths must match") else: try: other = type(self)._from_sequence(other)._data except (ValueError, TypeError): return ops.invalid_comparison(self, other, op) result = meth(self, other) result = com.values_from_object(result) o_mask = np.array(isna(other)) if o_mask.any(): result[o_mask] = nat_result if self._hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _period_array_cmp(cls, op): """ Wrap comparison operations to convert Period-like to PeriodDtype """ opname = '__{name}__'.format(name=op.__name__) nat_result = opname == '__ne__' def wrapper(self, other): op = getattr(self.asi8, opname) other = lib.item_from_zerodim(other) if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)): return NotImplemented if is_list_like(other) and len(other) != len(self): raise ValueError("Lengths must match") if isinstance(other, Period): self._check_compatible_with(other) result = op(other.ordinal) elif isinstance(other, cls): self._check_compatible_with(other) result = op(other.asi8) mask = self._isnan | other._isnan if mask.any(): result[mask] = nat_result return result elif other is NaT: result = np.empty(len(self.asi8), dtype=bool) result.fill(nat_result) else: other = Period(other, freq=self.freq) result = op(other.ordinal) if self._hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _create_comparison_method(cls, op): op_name = op.__name__ def cmp_method(self, other): if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)): # Rely on pandas to unbox and dispatch to us. return NotImplemented other = lib.item_from_zerodim(other) mask = None 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") # 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 else: mask = self._mask | mask result[mask] = op_name == "ne" return BooleanArray(result, np.zeros(len(result), dtype=bool), copy=False) name = "__{name}__".format(name=op.__name__) return set_function_name(cmp_method, name, cls)
def _td_index_cmp(opname, cls): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ nat_result = True if opname == '__ne__' else False def wrapper(self, other): msg = "cannot compare a TimedeltaIndex with type {0}" func = getattr(super(TimedeltaIndex, self), opname) if _is_convertible_to_td(other) or other is NaT: try: other = _to_m8(other) except ValueError: # failed to parse as timedelta raise TypeError(msg.format(type(other))) result = func(other) if isna(other): result.fill(nat_result) else: if not is_list_like(other): raise TypeError(msg.format(type(other))) other = TimedeltaIndex(other).values result = func(other) result = com._values_from_object(result) if isinstance(other, Index): o_mask = other.values.view('i8') == iNaT else: o_mask = other.view('i8') == iNaT if o_mask.any(): result[o_mask] = nat_result if self.hasnans: result[self._isnan] = nat_result # support of bool dtype indexers if is_bool_dtype(result): return result return Index(result) return compat.set_function_name(wrapper, opname, cls)
def _period_array_cmp(cls, op): """ Wrap comparison operations to convert Period-like to PeriodDtype """ opname = '__{name}__'.format(name=op.__name__) nat_result = True if opname == '__ne__' else False def wrapper(self, other): op = getattr(self.asi8, opname) if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)): return NotImplemented if is_list_like(other) and len(other) != len(self): raise ValueError("Lengths must match") if isinstance(other, Period): self._check_compatible_with(other) result = op(other.ordinal) elif isinstance(other, cls): self._check_compatible_with(other) result = op(other.asi8) mask = self._isnan | other._isnan if mask.any(): result[mask] = nat_result return result elif other is NaT: result = np.empty(len(self.asi8), dtype=bool) result.fill(nat_result) else: other = Period(other, freq=self.freq) result = op(other.ordinal) if self._hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _create_arithmetic_method(cls, op): def arithmetic_method(self, other): if isinstance(other, (ABCIndexClass, ABCSeries)): return NotImplemented elif isinstance(other, cls): other = other._ndarray with np.errstate(all="ignore"): result = op(self._ndarray, other) if op is divmod: a, b = result return cls(a), cls(b) return cls(result) return compat.set_function_name(arithmetic_method, f"__{op.__name__}__", cls)
def _td_array_cmp(cls, op): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ opname = '__{name}__'.format(name=op.__name__) nat_result = True if opname == '__ne__' else False meth = getattr(dtl.DatetimeLikeArrayMixin, opname) def wrapper(self, other): if _is_convertible_to_td(other) or other is NaT: try: other = _to_m8(other) except ValueError: # failed to parse as timedelta return ops.invalid_comparison(self, other, op) result = meth(self, other) if isna(other): result.fill(nat_result) elif not is_list_like(other): return ops.invalid_comparison(self, other, op) else: try: other = type(self)._from_sequence(other)._data except (ValueError, TypeError): return ops.invalid_comparison(self, other, op) result = meth(self, other) result = com.values_from_object(result) o_mask = np.array(isna(other)) if o_mask.any(): result[o_mask] = nat_result if self.hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _td_index_cmp(opname, cls, nat_result=False): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ def wrapper(self, other): msg = "cannot compare a TimedeltaIndex with type {0}" func = getattr(super(TimedeltaIndex, self), opname) if _is_convertible_to_td(other) or other is NaT: try: other = _to_m8(other) except ValueError: # failed to parse as timedelta raise TypeError(msg.format(type(other))) result = func(other) if isna(other): result.fill(nat_result) else: if not is_list_like(other): raise TypeError(msg.format(type(other))) other = TimedeltaIndex(other).values result = func(other) result = _values_from_object(result) if isinstance(other, Index): o_mask = other.values.view('i8') == iNaT else: o_mask = other.view('i8') == iNaT if o_mask.any(): result[o_mask] = nat_result if self.hasnans: result[self._isnan] = nat_result # support of bool dtype indexers if is_bool_dtype(result): return result return Index(result) return compat.set_function_name(wrapper, opname, cls)
def _create_arithmetic_method(cls, op): def arithmetic_method(self, other): if isinstance(other, (ABCIndexClass, ABCSeries)): return NotImplemented elif isinstance(other, cls): other = other._ndarray with np.errstate(all="ignore"): result = op(self._ndarray, other) if op is divmod: a, b = result return cls(a), cls(b) return cls(result) return compat.set_function_name(arithmetic_method, "__{}__".format(op.__name__), cls)
def _create_comparison_method(cls, op): op_name = op.__name__ if op_name in {"and_", "or_"}: op_name = op_name[:-1] def cmp_method(self, other): if isinstance(other, (ABCSeries, ABCIndexClass)): # Rely on pandas to unbox and dispatch to us. return NotImplemented if not is_scalar(other) and not isinstance(other, type(self)): # convert list-like to ndarray other = np.asarray(other) if isinstance(other, np.ndarray): # TODO: make this more flexible than just ndarray... if len(self) != len(other): raise AssertionError( "length mismatch: {self} vs. {other}".format( self=len(self), other=len(other) ) ) other = SparseArray(other, fill_value=self.fill_value) if isinstance(other, SparseArray): return _sparse_array_op(self, other, op, op_name) else: with np.errstate(all="ignore"): fill_value = op(self.fill_value, other) result = op(self.sp_values, other) return type(self)( result, sparse_index=self.sp_index, fill_value=fill_value, dtype=np.bool_, ) name = "__{name}__".format(name=op.__name__) return compat.set_function_name(cmp_method, name, cls)
def _td_index_cmp(opname, cls): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ nat_result = True if opname == '__ne__' else False def wrapper(self, other): msg = "cannot compare a {cls} with type {typ}" func = getattr(super(TimedeltaIndex, self), opname) if _is_convertible_to_td(other) or other is NaT: try: other = _to_m8(other) except ValueError: # failed to parse as timedelta raise TypeError(msg.format(cls=type(self).__name__, typ=type(other).__name__)) result = func(other) if isna(other): result.fill(nat_result) elif not is_list_like(other): raise TypeError(msg.format(cls=type(self).__name__, typ=type(other).__name__)) else: other = TimedeltaIndex(other).values result = func(other) result = com._values_from_object(result) o_mask = np.array(isna(other)) if o_mask.any(): result[o_mask] = nat_result if self.hasnans: result[self._isnan] = nat_result # support of bool dtype indexers if is_bool_dtype(result): return result return Index(result) return compat.set_function_name(wrapper, opname, cls)
def _create_arithmetic_method(cls, op): op_name = op.__name__ @unpack_zerodim_and_defer(op_name) def sparse_arithmetic_method(self, other): if isinstance(other, SparseArray): return _sparse_array_op(self, other, op, op_name) elif is_scalar(other): with np.errstate(all="ignore"): fill = op(_get_fill(self), np.asarray(other)) result = op(self.sp_values, other) if op_name == "divmod": left, right = result lfill, rfill = fill return ( _wrap_result(op_name, left, self.sp_index, lfill), _wrap_result(op_name, right, self.sp_index, rfill), ) return _wrap_result(op_name, result, self.sp_index, fill) else: other = np.asarray(other) with np.errstate(all="ignore"): # TODO: look into _wrap_result if len(self) != len(other): raise AssertionError( (f"length mismatch: {len(self)} vs. {len(other)}") ) if not isinstance(other, SparseArray): dtype = getattr(other, "dtype", None) other = SparseArray( other, fill_value=self.fill_value, dtype=dtype ) return _sparse_array_op(self, other, op, op_name) name = f"__{op.__name__}__" return compat.set_function_name(sparse_arithmetic_method, name, cls)
def _period_index_cmp(opname, cls, nat_result=False): """ Wrap comparison operations to convert datetime-like to datetime64 """ def wrapper(self, other): if isinstance(other, Period): func = getattr(self._ndarray_values, opname) other_base, _ = _gfc(other.freq) if other.freq != self.freq: msg = _DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr) raise IncompatibleFrequency(msg) result = func(other.ordinal) elif isinstance(other, PeriodIndex): if other.freq != self.freq: msg = _DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr) raise IncompatibleFrequency(msg) op = getattr(self._ndarray_values, opname) result = op(other._ndarray_values) mask = self._isnan | other._isnan if mask.any(): result[mask] = nat_result return result elif other is tslib.NaT: result = np.empty(len(self._ndarray_values), dtype=bool) result.fill(nat_result) else: other = Period(other, freq=self.freq) func = getattr(self._ndarray_values, opname) result = func(other.ordinal) if self.hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _create_arithmetic_method(cls, op): def sparse_arithmetic_method(self, other): op_name = op.__name__ if isinstance(other, (ABCSeries, ABCIndexClass)): other = getattr(other, 'values', other) if isinstance(other, SparseArray): return _sparse_array_op(self, other, op, op_name) elif is_scalar(other): with np.errstate(all='ignore'): fill = op(_get_fill(self), np.asarray(other)) result = op(self.sp_values, other) if op_name == 'divmod': left, right = result lfill, rfill = fill return (_wrap_result(op_name, left, self.sp_index, lfill), _wrap_result(op_name, right, self.sp_index, rfill)) return _wrap_result(op_name, result, self.sp_index, fill) else: other = np.asarray(other) with np.errstate(all='ignore'): # TODO: delete sparse stuff in core/ops.py # TODO: look into _wrap_result if len(self) != len(other): raise AssertionError( ("length mismatch: {self} vs. {other}".format( self=len(self), other=len(other)))) if not isinstance(other, SparseArray): dtype = getattr(other, 'dtype', None) other = SparseArray(other, fill_value=self.fill_value, dtype=dtype) return _sparse_array_op(self, other, op, op_name) name = '__{name}__'.format(name=op.__name__) return compat.set_function_name(sparse_arithmetic_method, name, cls)
def _create_logical_method(cls, op): @ops.unpack_zerodim_and_defer(op.__name__) def logical_method(self, other): assert op.__name__ in {"or_", "ror_", "and_", "rand_", "xor", "rxor"} other_is_booleanarray = isinstance(other, BooleanArray) other_is_scalar = lib.is_scalar(other) mask = None if other_is_booleanarray: other, mask = other._data, other._mask elif is_list_like(other): other = np.asarray(other, dtype="bool") if other.ndim > 1: raise NotImplementedError( "can only perform ops with 1-d structures" ) other, mask = coerce_to_array(other, copy=False) elif isinstance(other, np.bool_): other = other.item() if other_is_scalar and not (other is libmissing.NA or lib.is_bool(other)): raise TypeError( "'other' should be pandas.NA or a bool. " f"Got {type(other).__name__} instead." ) if not other_is_scalar and len(self) != len(other): raise ValueError("Lengths must match to compare") if op.__name__ in {"or_", "ror_"}: result, mask = ops.kleene_or(self._data, other, self._mask, mask) elif op.__name__ in {"and_", "rand_"}: result, mask = ops.kleene_and(self._data, other, self._mask, mask) elif op.__name__ in {"xor", "rxor"}: result, mask = ops.kleene_xor(self._data, other, self._mask, mask) return BooleanArray(result, mask) name = f"__{op.__name__}__" return set_function_name(logical_method, name, cls)
def _create_comparison_method(cls, op): op_name = op.__name__ @unpack_zerodim_and_defer(op.__name__) def cmp_method(self, other): mask = None if isinstance(other, IntegerArray): 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") # 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"): 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 else: mask = self._mask | mask result[mask] = op_name == "ne" return result name = f"__{op.__name__}__" return set_function_name(cmp_method, name, cls)
def _period_array_cmp(cls, op): """ Wrap comparison operations to convert Period-like to PeriodDtype """ opname = "__{name}__".format(name=op.__name__) nat_result = opname == "__ne__" @unpack_zerodim_and_defer(opname) def wrapper(self, other): ordinal_op = getattr(self.asi8, opname) if is_list_like(other) and len(other) != len(self): raise ValueError("Lengths must match") if isinstance(other, Period): self._check_compatible_with(other) result = ordinal_op(other.ordinal) elif isinstance(other, cls): self._check_compatible_with(other) result = ordinal_op(other.asi8) mask = self._isnan | other._isnan if mask.any(): result[mask] = nat_result return result elif other is NaT: result = np.empty(len(self.asi8), dtype=bool) result.fill(nat_result) else: other = Period(other, freq=self.freq) result = ordinal_op(other.ordinal) if self._hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _create_method(cls, op, coerce_to_dtype=True, result_dtype=None): """ Add support for binary operators by unwrapping, applying, and rewrapping. """ # NOTE(Clark): This overrides, but coerce_to_dtype, result_dtype might # not be needed def _binop(self, other): lvalues = self._tensor if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndex)): # Rely on Pandas to unbox and dispatch to us. return NotImplemented # divmod returns a tuple if op_name in ["__divmod__", "__rdivmod__"]: # TODO(Clark): Add support for divmod and rdivmod. # div, mod = result raise NotImplementedError if isinstance(other, (TensorArray, TensorArrayElement)): rvalues = other._tensor else: rvalues = other result = op(lvalues, rvalues) # Force a TensorArray if rvalue is not a scalar. if (isinstance(self, TensorArrayElement) and (not isinstance(other, TensorArrayElement) or not np.isscalar(other))): result_wrapped = TensorArray(result) else: result_wrapped = cls(result) return result_wrapped op_name = f"__{op.__name__}__" return set_function_name(_binop, op_name, cls)
def _period_array_cmp(cls, op): """ Wrap comparison operations to convert Period-like to PeriodDtype """ opname = '__{name}__'.format(name=op.__name__) nat_result = True if opname == '__ne__' else False def wrapper(self, other): op = getattr(self._ndarray_values, opname) if isinstance(other, Period): if other.freq != self.freq: msg = DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr) raise IncompatibleFrequency(msg) result = op(other.ordinal) elif isinstance(other, PeriodArrayMixin): if other.freq != self.freq: msg = DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr) raise IncompatibleFrequency(msg) result = op(other._ndarray_values) mask = self._isnan | other._isnan if mask.any(): result[mask] = nat_result return result elif other is NaT: result = np.empty(len(self._ndarray_values), dtype=bool) result.fill(nat_result) else: other = Period(other, freq=self.freq) result = op(other.ordinal) if self.hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _td_array_cmp(cls, op): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ opname = '__{name}__'.format(name=op.__name__) nat_result = True if opname == '__ne__' else False def wrapper(self, other): msg = "cannot compare a {cls} with type {typ}" meth = getattr(dtl.DatetimeLikeArrayMixin, opname) if _is_convertible_to_td(other) or other is NaT: try: other = _to_m8(other) except ValueError: # failed to parse as timedelta raise TypeError( msg.format(cls=type(self).__name__, typ=type(other).__name__)) result = meth(self, other) if isna(other): result.fill(nat_result) elif not is_list_like(other): raise TypeError( msg.format(cls=type(self).__name__, typ=type(other).__name__)) else: other = type(self)(other).values result = meth(self, other) result = com._values_from_object(result) o_mask = np.array(isna(other)) if o_mask.any(): result[o_mask] = nat_result if self.hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _td_array_cmp(cls, op): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ opname = '__{name}__'.format(name=op.__name__) nat_result = True if opname == '__ne__' else False def wrapper(self, other): msg = "cannot compare a {cls} with type {typ}" meth = getattr(dtl.DatetimeLikeArrayMixin, opname) if _is_convertible_to_td(other) or other is NaT: try: other = _to_m8(other) except ValueError: # failed to parse as timedelta raise TypeError(msg.format(cls=type(self).__name__, typ=type(other).__name__)) result = meth(self, other) if isna(other): result.fill(nat_result) elif not is_list_like(other): raise TypeError(msg.format(cls=type(self).__name__, typ=type(other).__name__)) else: other = type(self)(other)._data result = meth(self, other) result = com.values_from_object(result) o_mask = np.array(isna(other)) if o_mask.any(): result[o_mask] = nat_result if self.hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _create_comparison_method(cls, op): def cmp_method(self, other): op_name = op.__name__ mask = None if isinstance(other, (ABCSeries, ABCIndexClass)): # Rely on pandas to unbox and dispatch to us. return NotImplemented if isinstance(other, IntegerArray): other, mask = other._data, other._mask elif is_list_like(other): other = np.asarray(other) if other.ndim > 0 and len(self) != len(other): raise ValueError('Lengths must match to compare') other = lib.item_from_zerodim(other) # 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 else: mask = self._mask | mask result[mask] = op_name == 'ne' return result name = '__{name}__'.format(name=op.__name__) return set_function_name(cmp_method, name, cls)
def _create_comparison_method(cls, op): def cmp_method(self, other): op_name = op.__name__ mask = None if isinstance(other, (ABCSeries, ABCIndexClass)): # Rely on pandas to unbox and dispatch to us. return NotImplemented if isinstance(other, IntegerArray): other, mask = other._data, other._mask elif is_list_like(other): other = np.asarray(other) if other.ndim > 0 and len(self) != len(other): raise ValueError('Lengths must match to compare') other = lib.item_from_zerodim(other) # 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 else: mask = self._mask | mask result[mask] = True if op_name == 'ne' else False return result name = '__{name}__'.format(name=op.__name__) return set_function_name(cmp_method, name, cls)
def _create_arithmetic_method(cls, op): # Note: this handles both arithmetic and comparison methods. def method(self, other): from pandas.arrays import BooleanArray assert op.__name__ in ops.ARITHMETIC_BINOPS | ops.COMPARISON_BINOPS if isinstance(other, (ABCIndexClass, ABCSeries, ABCDataFrame)): return NotImplemented elif isinstance(other, cls): other = other._ndarray mask = isna(self) | isna(other) valid = ~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)}" ) other = np.asarray(other) other = other[valid] 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) return compat.set_function_name(method, f"__{op.__name__}__", cls)
def _create_comparison_method(cls, op): def cmp_method(self, other): op_name = op.__name__ if op_name in {'and_', 'or_'}: op_name = op_name[:-1] if isinstance(other, (ABCSeries, ABCIndexClass)): other = getattr(other, 'values', other) if not is_scalar(other) and not isinstance(other, type(self)): # convert list-like to ndarary other = np.asarray(other) if isinstance(other, np.ndarray): # TODO: make this more flexible than just ndarray... if len(self) != len(other): raise AssertionError( "length mismatch: {self} vs. {other}".format( self=len(self), other=len(other))) other = SparseArray(other, fill_value=self.fill_value) if isinstance(other, SparseArray): return _sparse_array_op(self, other, op, op_name) else: with np.errstate(all='ignore'): fill_value = op(self.fill_value, other) result = op(self.sp_values, other) return type(self)(result, sparse_index=self.sp_index, fill_value=fill_value, dtype=np.bool_) name = '__{name}__'.format(name=op.__name__) return compat.set_function_name(cmp_method, name, cls)
def _create_comparison_method(cls, op): def cmp_method(self, other): op_name = op.__name__ if op_name in {'and_', 'or_'}: op_name = op_name[:-1] if isinstance(other, (ABCSeries, ABCIndexClass)): other = getattr(other, 'values', other) if not is_scalar(other) and not isinstance(other, type(self)): # convert list-like to ndarary other = np.asarray(other) if isinstance(other, np.ndarray): # TODO: make this more flexible than just ndarray... if len(self) != len(other): raise AssertionError("length mismatch: {self} vs. {other}" .format(self=len(self), other=len(other))) other = SparseArray(other, fill_value=self.fill_value) if isinstance(other, SparseArray): return _sparse_array_op(self, other, op, op_name) else: with np.errstate(all='ignore'): fill_value = op(self.fill_value, other) result = op(self.sp_values, other) return type(self)(result, sparse_index=self.sp_index, fill_value=fill_value, dtype=np.bool_) name = '__{name}__'.format(name=op.__name__) return compat.set_function_name(cmp_method, name, cls)
def _create_comparison_method(cls, op): op_name = op.__name__ if op_name in {"and_", "or_"}: op_name = op_name[:-1] @unpack_zerodim_and_defer(op_name) def cmp_method(self, other): if not is_scalar(other) and not isinstance(other, type(self)): # convert list-like to ndarray other = np.asarray(other) if isinstance(other, np.ndarray): # TODO: make this more flexible than just ndarray... if len(self) != len(other): raise AssertionError( f"length mismatch: {len(self)} vs. {len(other)}" ) other = SparseArray(other, fill_value=self.fill_value) if isinstance(other, SparseArray): return _sparse_array_op(self, other, op, op_name) else: with np.errstate(all="ignore"): fill_value = op(self.fill_value, other) result = op(self.sp_values, other) return type(self)( result, sparse_index=self.sp_index, fill_value=fill_value, dtype=np.bool_, ) name = f"__{op.__name__}__" return compat.set_function_name(cmp_method, name, cls)
def _create_arithmetic_method(cls, op): op_name = op.__name__ @unpack_zerodim_and_defer(op.__name__) def integer_arithmetic_method(self, other): omask = None if getattr(other, "ndim", 0) > 1: raise NotImplementedError("can only perform ops with 1-d structures") if isinstance(other, IntegerArray): other, omask = 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") if not (is_float_dtype(other) or is_integer_dtype(other)): raise TypeError("can only perform ops with numeric values") else: if not (is_float(other) or is_integer(other) or other is libmissing.NA): raise TypeError("can only perform ops with numeric values") if omask is None: mask = self._mask.copy() if other is libmissing.NA: mask |= True else: mask = self._mask | omask if op_name == "pow": # 1 ** x is 1. mask = np.where((self._data == 1) & ~self._mask, False, mask) # x ** 0 is 1. if omask is not None: mask = np.where((other == 0) & ~omask, False, mask) elif other is not libmissing.NA: mask = np.where(other == 0, False, mask) elif op_name == "rpow": # 1 ** x is 1. if omask is not None: mask = np.where((other == 1) & ~omask, False, mask) elif other is not libmissing.NA: mask = np.where(other == 1, False, mask) # x ** 0 is 1. mask = np.where((self._data == 0) & ~self._mask, False, mask) if other is libmissing.NA: result = np.ones_like(self._data) else: with np.errstate(all="ignore"): result = op(self._data, other) # divmod returns a tuple if op_name == "divmod": div, mod = result return ( self._maybe_mask_result(div, mask, other, "floordiv"), self._maybe_mask_result(mod, mask, other, "mod"), ) return self._maybe_mask_result(result, mask, other, op_name) name = f"__{op.__name__}__" return set_function_name(integer_arithmetic_method, name, cls)
def _make_evaluate_binop(op, step=False): """ Parameters ---------- op : callable that accepts 2 params perform the binary op step : callable, optional, default to False op to apply to the step parm if not None if False, use the existing step """ @unpack_zerodim_and_defer(op.__name__) def _evaluate_numeric_binop(self, other): if isinstance(other, ABCTimedeltaIndex): # Defer to TimedeltaIndex implementation return NotImplemented elif isinstance(other, (timedelta, np.timedelta64)): # GH#19333 is_integer evaluated True on timedelta64, # so we need to catch these explicitly return op(self._int64index, other) elif is_timedelta64_dtype(other): # Must be an np.ndarray; GH#22390 return op(self._int64index, other) other = extract_array(other, extract_numpy=True) attrs = self._get_attributes_dict() left, right = self, other try: # apply if we have an override if step: with np.errstate(all="ignore"): rstep = step(left.step, right) # we don't have a representable op # so return a base index if not is_integer(rstep) or not rstep: raise ValueError else: rstep = left.step with np.errstate(all="ignore"): rstart = op(left.start, right) rstop = op(left.stop, right) result = type(self)(rstart, rstop, rstep, **attrs) # for compat with numpy / Int64Index # even if we can represent as a RangeIndex, return # as a Float64Index if we have float-like descriptors if not all(is_integer(x) for x in [rstart, rstop, rstep]): result = result.astype("float64") return result except (ValueError, TypeError, ZeroDivisionError): # Defer to Int64Index implementation return op(self._int64index, other) # TODO: Do attrs get handled reliably? name = f"__{op.__name__}__" return compat.set_function_name(_evaluate_numeric_binop, name, cls)
def _td_array_cmp(cls, op): """ Wrap comparison operations to convert timedelta-like to timedelta64 """ opname = f"__{op.__name__}__" nat_result = opname == "__ne__" @unpack_zerodim_and_defer(opname) def wrapper(self, other): if isinstance(other, str): try: other = self._scalar_from_string(other) except ValueError: # failed to parse as timedelta return invalid_comparison(self, other, op) if isinstance(other, self._recognized_scalars) or other is NaT: other = self._scalar_type(other) self._check_compatible_with(other) other_i8 = self._unbox_scalar(other) result = op(self.view("i8"), other_i8) if isna(other): result.fill(nat_result) elif not is_list_like(other): return invalid_comparison(self, other, op) elif len(other) != len(self): raise ValueError("Lengths must match") else: if isinstance(other, list): other = np.array(other) if not isinstance(other, (np.ndarray, cls)): return invalid_comparison(self, other, op) if is_object_dtype(other): with np.errstate(all="ignore"): result = ops.comp_method_OBJECT_ARRAY( op, self.astype(object), other ) o_mask = isna(other) elif not cls._is_recognized_dtype(other.dtype): # e.g. other is datetimearray return invalid_comparison(self, other, op) else: other = type(self)._from_sequence(other) self._check_compatible_with(other) result = op(self.view("i8"), other.view("i8")) o_mask = other._isnan if o_mask.any(): result[o_mask] = nat_result if self._hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)
def _create_method(cls, op, coerce_to_dtype=True): """ A class method that returns a method that will correspond to an operator for an ExtensionArray subclass, by dispatching to the relevant operator defined on the individual elements of the ExtensionArray. Parameters ---------- op : function An operator that takes arguments op(a, b) coerce_to_dtype : bool boolean indicating whether to attempt to convert the result to the underlying ExtensionArray dtype (default True) Returns ------- A method that can be bound to a method of a class Example ------- Given an ExtensionArray subclass called MyExtensionArray, use >>> __add__ = cls._create_method(operator.add) in the class definition of MyExtensionArray to create the operator for addition, that will be based on the operator implementation of the underlying elements of the ExtensionArray """ def _binop(self, other): def validate_length(obj1, obj2): # validates length and converts to listlike try: if len(obj1) == len(obj2): return obj2 else: raise ValueError("Lengths must match") except TypeError: return [obj2] * len(obj1) def convert_values(param): # convert to a quantity or listlike if isinstance(param, cls): return param.quantity elif isinstance(param, _Quantity): return param elif is_list_like(param) and isinstance(param[0], _Quantity): return type(param[0])([p.magnitude for p in param], param[0].units) else: return param if isinstance(other, Series): return NotImplemented lvalues = self.quantity other = validate_length(lvalues, other) rvalues = convert_values(other) # Pint quantities may only be exponented by single values, not arrays. # Reduce single value arrays to single value to allow power ops if isinstance(rvalues, _Quantity): if len(set(np.array(rvalues.data))) == 1: rvalues = rvalues[0] elif len(set(np.array(rvalues))) == 1: rvalues = rvalues[0] # If the operator is not defined for the underlying objects, # a TypeError should be raised res = op(lvalues, rvalues) if op.__name__ == "divmod": return ( cls.from_1darray_quantity(res[0]), cls.from_1darray_quantity(res[1]), ) if coerce_to_dtype: try: res = cls.from_1darray_quantity(res) except TypeError: pass return res op_name = ops._get_op_name(op, True) return set_function_name(_binop, op_name, cls)
def _create_method(cls, op, coerce_to_dtype=True): """ A class method that returns a method that will correspond to an operator for an ExtensionArray subclass, by dispatching to the relevant operator defined on the individual elements of the ExtensionArray. Parameters ---------- op : function An operator that takes arguments op(a, b) coerce_to_dtype : bool boolean indicating whether to attempt to convert the result to the underlying ExtensionArray dtype (default True) Returns ------- A method that can be bound to a method of a class Example ------- Given an ExtensionArray subclass called MyExtensionArray, use >>> __add__ = cls._create_method(operator.add) in the class definition of MyExtensionArray to create the operator for addition, that will be based on the operator implementation of the underlying elements of the ExtensionArray """ def _binop(self, other): def validate_length(obj1, obj2): # validates length # CHANGED: do not convert to listlike (why should we? pint.Quantity is perfecty able to handle that...) try: if len(obj1) != len(obj2): raise ValueError("Lengths must match") except TypeError: pass def convert_values(param): # convert to a quantity or listlike if isinstance(param, cls): return param.quantity elif isinstance(param, (_Quantity, _Unit)): return param elif (is_list_like(param) and len(param) > 0 and isinstance(param[0], _Quantity)): units = param[0].units return type(param[0])([p.m_as(units) for p in param], units) else: return param if isinstance(other, (Series, DataFrame)): return NotImplemented lvalues = self.quantity validate_length(lvalues, other) rvalues = convert_values(other) # If the operator is not defined for the underlying objects, # a TypeError should be raised res = op(lvalues, rvalues) if op.__name__ == "divmod": return ( cls.from_1darray_quantity(res[0]), cls.from_1darray_quantity(res[1]), ) if coerce_to_dtype: try: res = cls.from_1darray_quantity(res) except TypeError: pass return res op_name = f"__{op}__" return set_function_name(_binop, op_name, cls)
def _make_compare(op): opname = f"__{op.__name__}__" _evaluate_compare = make_wrapped_comparison_op(opname) return compat.set_function_name(_evaluate_compare, opname, cls)
def _create_method(cls, op, coerce_to_dtype=True): """ A class method that returns a method that will correspond to an operator for an ExtensionArray subclass, by dispatching to the relevant operator defined on the individual elements of the ExtensionArray. Parameters ---------- op : function An operator that takes arguments op(a, b) coerce_to_dtype : bool, default True boolean indicating whether to attempt to convert the result to the underlying ExtensionArray dtype. If it's not possible to create a new ExtensionArray with the values, an ndarray is returned instead. Returns ------- Callable[[Any, Any], Union[ndarray, ExtensionArray]] A method that can be bound to a class. When used, the method receives the two arguments, one of which is the instance of this class, and should return an ExtensionArray or an ndarray. Returning an ndarray may be necessary when the result of the `op` cannot be stored in the ExtensionArray. The dtype of the ndarray uses NumPy's normal inference rules. Example ------- Given an ExtensionArray subclass called MyExtensionArray, use >>> __add__ = cls._create_method(operator.add) in the class definition of MyExtensionArray to create the operator for addition, that will be based on the operator implementation of the underlying elements of the ExtensionArray """ def _binop(self, other): def convert_values(param): if isinstance(param, ExtensionArray) or is_list_like(param): ovalues = param else: # Assume its an object ovalues = [param] * len(self) return ovalues lvalues = self rvalues = convert_values(other) # If the operator is not defined for the underlying objects, # a TypeError should be raised res = [op(a, b) for (a, b) in zip(lvalues, rvalues)] def _maybe_convert(arr): if coerce_to_dtype: # https://github.com/pandas-dev/pandas/issues/22850 # We catch all regular exceptions here, and fall back # to an ndarray. try: res = self._from_sequence(arr) except Exception: res = np.asarray(arr) else: res = np.asarray(arr) return res if op.__name__ in {'divmod', 'rdivmod'}: a, b = zip(*res) res = _maybe_convert(a), _maybe_convert(b) else: res = _maybe_convert(res) return res op_name = ops._get_op_name(op, True) return set_function_name(_binop, op_name, cls)
def _make_evaluate_binop(op, step=False): """ Parameters ---------- op : callable that accepts 2 parms perform the binary op step : callable, optional, default to False op to apply to the step parm if not None if False, use the existing step """ def _evaluate_numeric_binop(self, other): if isinstance(other, (ABCSeries, ABCDataFrame)): return NotImplemented elif isinstance(other, ABCTimedeltaIndex): # Defer to TimedeltaIndex implementation return NotImplemented elif isinstance(other, (timedelta, np.timedelta64)): # GH#19333 is_integer evaluated True on timedelta64, # so we need to catch these explicitly return op(self._int64index, other) elif is_timedelta64_dtype(other): # Must be an np.ndarray; GH#22390 return op(self._int64index, other) other = self._validate_for_numeric_binop(other, op) attrs = self._get_attributes_dict() attrs = self._maybe_update_attributes(attrs) left, right = self, other try: # apply if we have an override if step: with np.errstate(all='ignore'): rstep = step(left._step, right) # we don't have a representable op # so return a base index if not is_integer(rstep) or not rstep: raise ValueError else: rstep = left._step with np.errstate(all='ignore'): rstart = op(left._start, right) rstop = op(left._stop, right) result = RangeIndex(rstart, rstop, rstep, **attrs) # for compat with numpy / Int64Index # even if we can represent as a RangeIndex, return # as a Float64Index if we have float-like descriptors if not all(is_integer(x) for x in [rstart, rstop, rstep]): result = result.astype('float64') return result except (ValueError, TypeError, ZeroDivisionError): # Defer to Int64Index implementation return op(self._int64index, other) # TODO: Do attrs get handled reliably? name = '__{name}__'.format(name=op.__name__) return compat.set_function_name(_evaluate_numeric_binop, name, cls)
def _create_arithmetic_method(cls, op): # Note: this handles both arithmetic and comparison methods. def method(self, other): is_arithmetic = \ True if op.__name__ in ops.ARITHMETIC_BINOPS else False pandas_only = cls._pandas_only() is_other_array = False if not is_scalar(other): is_other_array = True other = np.asarray(other) self_is_na = self.isna() other_is_na = pd.isna(other) mask = self_is_na | other_is_na if pa is None or pandas_only: if is_arithmetic: ret = np.empty(self.shape, dtype=object) else: ret = np.zeros(self.shape, dtype=bool) valid = ~mask arr = self._arrow_array.to_pandas().to_numpy() \ if self._use_arrow else self._ndarray o = other[valid] if is_other_array else other ret[valid] = op(arr[valid], o) if is_arithmetic: return ArrowStringArray(ret) else: return pd.arrays.BooleanArray(ret, mask) chunks = [] mask_chunks = [] start = 0 for chunk_array in self._arrow_array.chunks: chunk_array = np.asarray(chunk_array.to_pandas()) end = start + len(chunk_array) chunk_mask = mask[start:end] chunk_valid = ~chunk_mask if is_arithmetic: result = np.empty(chunk_array.shape, dtype=object) else: result = np.zeros(chunk_array.shape, dtype=bool) chunk_other = other if is_other_array: chunk_other = other[start:end] chunk_other = chunk_other[chunk_valid] # calculate only for both not None result[chunk_valid] = op(chunk_array[chunk_valid], chunk_other) if is_arithmetic: chunks.append( pa.array(result, type=pa.string(), from_pandas=True)) else: chunks.append(result) mask_chunks.append(chunk_mask) if is_arithmetic: return ArrowStringArray(pa.chunked_array(chunks)) else: return pd.arrays.BooleanArray(np.concatenate(chunks), np.concatenate(mask_chunks)) return set_function_name(method, f"__{op.__name__}__", cls)
def _create_method(cls, op, coerce_to_dtype=True): """ A class method that returns a method that will correspond to an operator for an ExtensionArray subclass, by dispatching to the relevant operator defined on the individual elements of the ExtensionArray. Parameters ---------- op : function An operator that takes arguments op(a, b) coerce_to_dtype : bool, default True boolean indicating whether to attempt to convert the result to the underlying ExtensionArray dtype. If it's not possible to create a new ExtensionArray with the values, an ndarray is returned instead. Returns ------- Callable[[Any, Any], Union[ndarray, ExtensionArray]] A method that can be bound to a class. When used, the method receives the two arguments, one of which is the instance of this class, and should return an ExtensionArray or an ndarray. Returning an ndarray may be necessary when the result of the `op` cannot be stored in the ExtensionArray. The dtype of the ndarray uses NumPy's normal inference rules. Examples -------- Given an ExtensionArray subclass called MyExtensionArray, use >>> __add__ = cls._create_method(operator.add) in the class definition of MyExtensionArray to create the operator for addition, that will be based on the operator implementation of the underlying elements of the ExtensionArray """ def _binop(self, other): def convert_values(param): if isinstance(param, ExtensionArray) or is_list_like(param): ovalues = param else: # Assume its an object ovalues = [param] * len(self) return ovalues if isinstance(other, (ABCSeries, ABCIndexClass)): # rely on qq_pandas to unbox and dispatch to us return NotImplemented lvalues = self rvalues = convert_values(other) # If the operator is not defined for the underlying objects, # a TypeError should be raised res = [op(a, b) for (a, b) in zip(lvalues, rvalues)] def _maybe_convert(arr): if coerce_to_dtype: # https://github.com/pandas-dev/pandas/issues/22850 # We catch all regular exceptions here, and fall back # to an ndarray. try: res = self._from_sequence(arr) except Exception: res = np.asarray(arr) else: res = np.asarray(arr) return res if op.__name__ in {"divmod", "rdivmod"}: a, b = zip(*res) res = _maybe_convert(a), _maybe_convert(b) else: res = _maybe_convert(res) return res op_name = ops._get_op_name(op, True) return set_function_name(_binop, op_name, cls)
def _period_array_cmp(cls, op): """ Wrap comparison operations to convert Period-like to PeriodDtype """ opname = f"__{op.__name__}__" nat_result = opname == "__ne__" @unpack_zerodim_and_defer(opname) def wrapper(self, other): ordinal_op = getattr(self.asi8, opname) if is_list_like(other) and len(other) != len(self): raise ValueError("Lengths must match") if isinstance(other, str): try: other = self._scalar_from_string(other) except ValueError: # string that can't be parsed as Period return invalid_comparison(self, other, op) elif isinstance(other, int): # TODO: sure we want to allow this? we dont for DTA/TDA # 2 tests rely on this other = Period(other, freq=self.freq) result = ordinal_op(other.ordinal) if isinstance(other, Period): self._check_compatible_with(other) result = ordinal_op(other.ordinal) elif other is NaT: result = np.empty(len(self.asi8), dtype=bool) result.fill(nat_result) elif not is_list_like(other): return invalid_comparison(self, other, op) else: if isinstance(other, list): # TODO: could use pd.Index to do inference? other = np.array(other) if not isinstance(other, (np.ndarray, cls)): return invalid_comparison(self, other, op) if is_object_dtype(other): with np.errstate(all="ignore"): result = ops.comp_method_OBJECT_ARRAY( op, self.astype(object), other ) o_mask = isna(other) elif not is_period_dtype(other): # e.g. is_timedelta64_dtype(other) return invalid_comparison(self, other, op) else: assert isinstance(other, cls), type(other) self._check_compatible_with(other) result = ordinal_op(other.asi8) o_mask = other._isnan if o_mask.any(): result[o_mask] = nat_result if self._hasnans: result[self._isnan] = nat_result return result return compat.set_function_name(wrapper, opname, cls)