def _create_comparison_method(cls, op):
op_name = op.__name__
#print(op_name)
def cmp_method(self, other):
if is_list_like(other):
other = cls._from_sequence(other)
if other.ndim > 0 and len(self) != len(other):
raise ValueError("Lengths must match to compare")
if op_name == 'eq':
return SparseExtensionArray(self.data == other.data)
if op_name == 'ne':
return SparseExtensionArray(self.data != other.data)
if op_name == 'lt':
return SparseExtensionArray(self.data < other.data)
if op_name == 'gt':
return SparseExtensionArray(self.data > other.data)
if op_name == 'le':
return SparseExtensionArray(self.data <= other.data)
if op_name == 'ge':
return SparseExtensionArray(self.data >= other.data)
opname = ops._get_op_name(op, True)
return set_function_name(cmp_method, opname, cls)
def _create_method(cls, op, coerce_to_dtype=True, result_dtype=None):
# NOTE: this overrides, but coerce_to_dtype, result_dtype might not be needed
def _binop(self, other):
lvalues = self._tensor
if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
# Rely on pandas to unbox and dispatch to us.
return NotImplemented
# divmod returns a tuple
if op_name in ["__divmod__", "__rdivmod__"]:
# TODO: return tuple
# div, mod = result
raise NotImplementedError
if isinstance(other, (TensorArray, TensorElement)):
rvalues = other._tensor
else:
rvalues = other
result = op(lvalues, rvalues)
# Force a TensorArray if rvalue is not a scalar
if isinstance(self, TensorElement) and \
(not isinstance(other, TensorElement) or not np.isscalar(other)):
result_wrapped = TensorArray(result)
else:
result_wrapped = cls(result)
return result_wrapped
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 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 _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 _create_method(cls, op, coerce_to_dtype=True):
# NOTE: this overrides, but coerce_to_dtype might not be needed
def _binop(self, other):
lvalues = self._tensor
rvalues = other._tensor if isinstance(other,
TensorArray) else other
res = op(lvalues, rvalues)
return TensorArray(res)
op_name = ops._get_op_name(op, True)
return set_function_name(_binop, op_name, cls)
def _create_arithmetic_method(cls, op):
#print('op name: %s' % op.__name__)
def sparse_arithmetic_method(self, other):
op_name = op.__name__
if op_name == 'add':
coords, data = _sum(self.data.coords[0], self.data.data,
other.values.data.coords[0],
other.values.data.data)
coo = sparse.COO(coords, data, shape=self.shape)
return SparseExtensionArray(coo)
opname = ops._get_op_name(op, True)
return set_function_name(sparse_arithmetic_method, 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.
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 _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, 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)
