ENH: Implement EA unary ops

doc/source/whatsnew/v0.25.0.rst

@@ -135,6 +135,22 @@ the output will truncate, if it's wider than :attr:`options.display.width`
 (default: 80 characters).
 
 
+.. _whatsnew_0250.enhancements.extension_array_operators:
+
+``ExtensionArray`` operator support
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+A ``Series`` based on an ``ExtensionArray`` now supports arithmetic, comparison and unary
+operators (:issue:`19577`, :issue:`23313`). There are two approaches for providing operator support for an ``ExtensionArray``:
+
+1. Define each of the operators on your ``ExtensionArray`` subclass.
+2. Use an operator implementation from pandas that depends on operators that are already defined
+ on the underlying elements (scalars) of the ``ExtensionArray``.
+
+See the :ref:`ExtensionArray Operator Support
+<extending.extension.operator>` documentation section for details on both
+ways of adding operator support.
+
 .. _whatsnew_0250.enhancements.other:
 
 Other enhancements

pandas/conftest.py

@@ -223,6 +223,20 @@ def all_compare_operators(request):
  return request.param
 
 
+@pytest.fixture(params=['__pos__', '__neg__', '__inv__', '__invert__',
+ '__abs__'])
+def all_unary_operators(request):
+ """
+ Fixture for dunder names for common unary operations
+
+ * +
+ * -
+ * ~
+ * abs
+ """
+ return request.param
+
+
 @pytest.fixture(params=['__le__', '__lt__', '__ge__', '__gt__'])
 def compare_operators_no_eq_ne(request):
  """
@@ -413,6 +427,7 @@ def tz_aware_fixture(request):
 SIGNED_EA_INT_DTYPES = ["Int8", "Int16", "Int32", "Int64"]
 ALL_INT_DTYPES = UNSIGNED_INT_DTYPES + SIGNED_INT_DTYPES
 ALL_EA_INT_DTYPES = UNSIGNED_EA_INT_DTYPES + SIGNED_EA_INT_DTYPES
+ALL_NUMPY_EA_INT_DTYPES = ALL_INT_DTYPES + ALL_EA_INT_DTYPES
 
 FLOAT_DTYPES = [float, "float32", "float64"]
 COMPLEX_DTYPES = [complex, "complex64", "complex128"]
@@ -556,6 +571,51 @@ def any_int_dtype(request):
  return request.param
 
 
+@pytest.fixture(params=ALL_EA_INT_DTYPES)
+def any_ea_int_dtype(request):
+ """
+ Parameterized fixture for any integer dtype.
+
+ * 'Int8'
+ * 'UInt8'
+ * 'Int16'
+ * 'UInt16'
+ * 'Int32'
+ * 'UInt32'
+ * 'Int64'
+ * 'UInt64'
+ """
+
+ return request.param
+
+
+@pytest.fixture(params=ALL_NUMPY_EA_INT_DTYPES)
+def any_numpy_ea_int_dtype(request):
+ """
+ Parameterized fixture for any integer dtype.
+
+ * int
+ * 'int8'
+ * 'uint8'
+ * 'int16'
+ * 'uint16'
+ * 'int32'
+ * 'uint32'
+ * 'int64'
+ * 'uint64'
+ * 'Int8'
+ * 'UInt8'
+ * 'Int16'
+ * 'UInt16'
+ * 'Int32'
+ * 'UInt32'
+ * 'Int64'
+ * 'UInt64'
+ """
+
+ return request.param
+
+
 @pytest.fixture(params=ALL_REAL_DTYPES)
 def any_real_dtype(request):
  """

pandas/core/arrays/base.py

@@ -1001,6 +1001,14 @@ def _add_comparison_ops(cls):
  cls.__le__ = cls._create_comparison_method(operator.le)
  cls.__ge__ = cls._create_comparison_method(operator.ge)
 
+ @classmethod
+ def _add_unary_ops(cls):
+ cls.__pos__ = cls._create_unary_method(operator.pos)
+ cls.__neg__ = cls._create_unary_method(operator.neg)
+ cls.__inv__ = cls._create_unary_method(operator.invert)
+ cls.__invert__ = cls._create_unary_method(operator.invert)
+ cls.__abs__ = cls._create_unary_method(operator.abs)
+
 
 class ExtensionScalarOpsMixin(ExtensionOpsMixin):
  """
@@ -1039,7 +1047,7 @@ def _create_method(cls, op, coerce_to_dtype=True):
  Parameters
  ----------
  op : function
- An operator that takes arguments op(a, b)
+ An operator that takes binary arguments op(a, b)
  coerce_to_dtype : bool, default True
  boolean indicating whether to attempt to convert
  the result to the underlying ExtensionArray dtype.
@@ -1087,24 +1095,12 @@ def convert_values(param):
  # 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)
+ res = (self._maybe_convert(a, coerce_to_dtype),
+ self._maybe_convert(b, coerce_to_dtype))
  else:
- res = _maybe_convert(res)
+ res = self._maybe_convert(res, coerce_to_dtype)
  return res
 
  op_name = ops._get_op_name(op, True)
@@ -1117,3 +1113,65 @@ def _create_arithmetic_method(cls, op):
  @classmethod
  def _create_comparison_method(cls, op):
  return cls._create_method(op, coerce_to_dtype=False)
+
+ @classmethod
+ def _create_unary_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 unary argument op(a)
+ 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], Union[ndarray, ExtensionArray]]
+ A method that can be bound to a class. When used, the method
+ receives 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
+ >>> __neg__ = cls._create_method(operator.neg)
+
+ in the class definition of MyExtensionArray to create the operator
+ for negative, that will be based on the operator implementation
+ of the underlying elements of the ExtensionArray
+ """
+
+ def _unaryop(self):
+ # If the operator is not defined for the underlying objects,
+ # a TypeError should be raised
+ res = [op(a) for a in self]
+ res = self._maybe_convert(res, coerce_to_dtype)
+ return res
+
+ op_name = ops._get_op_name(op, True)
+ return set_function_name(_unaryop, op_name, cls)
+
+ def _maybe_convert(self, arr, coerce_to_dtype):
+ 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

pandas/core/arrays/datetimes.py

@@ -785,6 +785,10 @@ def _add_delta(self, delta):
  new_values = super()._add_delta(delta)
  return type(self)._from_sequence(new_values, tz=self.tz, freq='infer')
 
+ def __pos__(self):
+ raise TypeError("Unary plus expects numeric dtype, not {}"
+ .format(self.dtype))
+
  # -----------------------------------------------------------------
  # Timezone Conversion and Localization Methods
 

pandas/core/arrays/integer.py

@@ -671,9 +671,20 @@ def integer_arithmetic_method(self, other):
  name = '__{name}__'.format(name=op.__name__)
  return set_function_name(integer_arithmetic_method, name, cls)
 
+ @classmethod
+ def _create_unary_method(cls, op):
+ def integer_unary_method(self):
+ with np.errstate(all='ignore'):
+ result = op(self._data)
+ return type(self)(result, self._mask)
+
+ name = '__{name}__'.format(name=op.__name__)
+ return set_function_name(integer_unary_method, name, cls)
+
 
 IntegerArray._add_arithmetic_ops()
 IntegerArray._add_comparison_ops()
+IntegerArray._add_unary_ops()
 
 
 module = sys.modules[__name__]

pandas/core/arrays/sparse.py

@@ -1723,9 +1723,6 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
  else:
  return type(self)(result)
 
- def __abs__(self):
- return np.abs(self)
-
  # ------------------------------------------------------------------------
  # Ops
  # ------------------------------------------------------------------------
@@ -1828,6 +1825,7 @@ def _add_unary_ops(cls):
  cls.__pos__ = cls._create_unary_method(operator.pos)
  cls.__neg__ = cls._create_unary_method(operator.neg)
  cls.__invert__ = cls._create_unary_method(operator.invert)
+ cls.__abs__ = cls._create_unary_method(operator.abs)
 
  @classmethod
  def _add_comparison_ops(cls):

pandas/core/arrays/timedeltas.py

@@ -733,6 +733,9 @@ def __rdivmod__(self, other):
  return res1, res2
 
  # Note: TimedeltaIndex overrides this in call to cls._add_numeric_methods
+ def __pos__(self):
+ return self.copy()
+
  def __neg__(self):
  if self.freq is not None:
  return type(self)(-self._data, freq=-self.freq)

pandas/core/frame.py

@@ -3302,7 +3302,7 @@ def _setitem_frame(self, key, value):
 
  self._check_inplace_setting(value)
  self._check_setitem_copy()
- self._where(-key, value, inplace=True)
+ self._where(~key, value, inplace=True)
 
  def _ensure_valid_index(self, value):
  """
@@ -4587,11 +4587,11 @@ def drop_duplicates(self, subset=None, keep='first', inplace=False):
  duplicated = self.duplicated(subset, keep=keep)
 
  if inplace:
- inds, = (-duplicated)._ndarray_values.nonzero()
+ inds, = (~duplicated)._ndarray_values.nonzero()
  new_data = self._data.take(inds)
  self._update_inplace(new_data)
  else:
- return self[-duplicated]
+ return self[~duplicated]
 
  def duplicated(self, subset=None, keep='first'):
  """