diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 8193d65b3b30c..9e2f1607bda9d 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -1171,6 +1171,15 @@ class ExtensionOpsMixin: def _create_arithmetic_method(cls, op): raise AbstractMethodError(cls) + @classmethod + def _create_unary_method(cls, op): + raise AbstractMethodError(cls) + + @classmethod + def _add_unary_ops(cls): + cls.__pos__ = cls._create_unary_method(operator.pos) + cls.__neg__ = cls._create_unary_method(operator.neg) + @classmethod def _add_arithmetic_ops(cls): cls.__add__ = cls._create_arithmetic_method(operator.add) @@ -1244,7 +1253,7 @@ class ExtensionScalarOpsMixin(ExtensionOpsMixin): """ @classmethod - def _create_method(cls, op, coerce_to_dtype=True, result_dtype=None): + def _create_method(cls, op, coerce_to_dtype=True, result_dtype=None, unary=False): """ A class method that returns a method that will correspond to an operator for an ExtensionArray subclass, by dispatching to the @@ -1283,6 +1292,24 @@ def _create_method(cls, op, coerce_to_dtype=True, result_dtype=None): of the underlying elements of the ExtensionArray """ + def _maybe_convert(self, 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. + res = maybe_cast_to_extension_array(type(self), arr) + if not isinstance(res, type(self)): + # exception raised in _from_sequence; ensure we have ndarray + res = np.asarray(arr) + else: + res = np.asarray(arr, dtype=result_dtype) + return res + + def _unaryop(self): + res = [op(a) for a in self] + + return _maybe_convert(self, res) + def _binop(self, other): def convert_values(param): if isinstance(param, ExtensionArray) or is_list_like(param): @@ -1302,26 +1329,15 @@ 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. - res = maybe_cast_to_extension_array(type(self), arr) - if not isinstance(res, type(self)): - # exception raised in _from_sequence; ensure we have ndarray - res = np.asarray(arr) - else: - res = np.asarray(arr, dtype=result_dtype) - return res - if op.__name__ in {"divmod", "rdivmod"}: a, b = zip(*res) - return _maybe_convert(a), _maybe_convert(b) + return _maybe_convert(self, a), _maybe_convert(self, b) - return _maybe_convert(res) + return _maybe_convert(self, res) op_name = f"__{op.__name__}__" + if unary: + return set_function_name(_unaryop, op_name, cls) return set_function_name(_binop, op_name, cls) @classmethod @@ -1331,3 +1347,7 @@ def _create_arithmetic_method(cls, op): @classmethod def _create_comparison_method(cls, op): return cls._create_method(op, coerce_to_dtype=False, result_dtype=bool) + + @classmethod + def _create_unary_method(cls, op): + return cls._create_method(op, unary=True) diff --git a/pandas/core/arrays/integer.py b/pandas/core/arrays/integer.py index d83ff91a1315f..8450dd50ba461 100644 --- a/pandas/core/arrays/integer.py +++ b/pandas/core/arrays/integer.py @@ -654,9 +654,24 @@ def integer_arithmetic_method(self, other): name = f"__{op.__name__}__" return set_function_name(integer_arithmetic_method, name, cls) + @classmethod + def _create_unary_method(cls, op): + op_name = op.__name__ + + @unpack_zerodim_and_defer(op.__name__) + def integer_unary_method(self): + mask = self._mask + with np.errstate(all="ignore"): + result = op(self._data) + return self._maybe_mask_result(result, mask, None, op_name) + + name = f"__{op.__name__}__" + return set_function_name(integer_unary_method, name, cls) + IntegerArray._add_arithmetic_ops() IntegerArray._add_comparison_ops() +IntegerArray._add_unary_ops() _dtype_docstring = """ diff --git a/pandas/core/ops/common.py b/pandas/core/ops/common.py index 515a0a5198d74..349e699426cf0 100644 --- a/pandas/core/ops/common.py +++ b/pandas/core/ops/common.py @@ -46,9 +46,13 @@ def _unpack_zerodim_and_defer(method, name: str): method """ is_cmp = name.strip("__") in {"eq", "ne", "lt", "le", "gt", "ge"} + is_unary = name.strip("__") in {"neg", "pos"} @wraps(method) - def new_method(self, other): + def new_method(self, other=None): + + if is_unary: + return method(self) if is_cmp and isinstance(self, ABCIndexClass) and isinstance(other, ABCSeries): # For comparison ops, Index does *not* defer to Series diff --git a/pandas/tests/extension/base/ops.py b/pandas/tests/extension/base/ops.py index c93603398977e..bbcace9e3e096 100644 --- a/pandas/tests/extension/base/ops.py +++ b/pandas/tests/extension/base/ops.py @@ -186,3 +186,15 @@ def test_invert(self, data): result = ~s expected = pd.Series(~data, name="name") self.assert_series_equal(result, expected) + + def test_neg(self, data): + s = pd.Series(data, name="name") + result = -s + expected = pd.Series(-data, name="name") + self.assert_series_equal(result, expected) + + def test_pos(self, data): + s = pd.Series(data, name="name") + result = +s + expected = pd.Series(+data, name="name") + self.assert_series_equal(result, expected) diff --git a/pandas/tests/extension/test_integer.py b/pandas/tests/extension/test_integer.py index 725533765ca2c..71b3c159169c1 100644 --- a/pandas/tests/extension/test_integer.py +++ b/pandas/tests/extension/test_integer.py @@ -181,6 +181,10 @@ def _compare_other(self, s, data, op_name, other): self.check_opname(s, op_name, other) +class TestUnaryOps(base.BaseUnaryOpsTests): + pass + + class TestInterface(base.BaseInterfaceTests): pass