From 3dadf1346f8ed493bdc3e031bea49f80da569df0 Mon Sep 17 00:00:00 2001 From: Asish Mahapatra Date: Wed, 2 Sep 2020 13:17:00 -0400 Subject: [PATCH 1/5] add unary ops ea support --- pandas/core/arrays/base.py | 34 +++++++++++++++++++++++++++++++++- pandas/core/arrays/integer.py | 15 +++++++++++++++ pandas/core/ops/common.py | 5 ++++- 3 files changed, 52 insertions(+), 2 deletions(-) diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 8193d65b3b30c..7371daff7879f 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,23 @@ def _create_method(cls, op, coerce_to_dtype=True, result_dtype=None): of the underlying elements of the ExtensionArray """ + def _unaryop(self): + res = [op(a) for a in self] + 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 + + return _maybe_convert(res) + def _binop(self, other): def convert_values(param): if isinstance(param, ExtensionArray) or is_list_like(param): @@ -1322,6 +1348,8 @@ def _maybe_convert(arr): return _maybe_convert(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 +1359,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..c81b6ae77141a 100644 --- a/pandas/core/arrays/integer.py +++ b/pandas/core/arrays/integer.py @@ -653,10 +653,25 @@ 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 unary_arithmetic_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(unary_arithmetic_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..ce03e7d333ddf 100644 --- a/pandas/core/ops/common.py +++ b/pandas/core/ops/common.py @@ -48,7 +48,10 @@ def _unpack_zerodim_and_defer(method, name: str): is_cmp = name.strip("__") in {"eq", "ne", "lt", "le", "gt", "ge"} @wraps(method) - def new_method(self, other): + def new_method(self, other=None): + + if other is None: + return method(self) if is_cmp and isinstance(self, ABCIndexClass) and isinstance(other, ABCSeries): # For comparison ops, Index does *not* defer to Series From 380e6d4a8f029578d3a26fb66bed8582b6eeeac1 Mon Sep 17 00:00:00 2001 From: Asish Mahapatra Date: Wed, 2 Sep 2020 17:02:10 -0400 Subject: [PATCH 2/5] merge changes local --- pandas/core/arrays/base.py | 46 ++++++++++++++------------------------ pandas/core/ops/common.py | 3 ++- 2 files changed, 19 insertions(+), 30 deletions(-) diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 7371daff7879f..9e2f1607bda9d 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -1292,22 +1292,23 @@ def _create_method(cls, op, coerce_to_dtype=True, result_dtype=None, unary=False 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] - 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 - - return _maybe_convert(res) + + return _maybe_convert(self, res) def _binop(self, other): def convert_values(param): @@ -1328,24 +1329,11 @@ 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: diff --git a/pandas/core/ops/common.py b/pandas/core/ops/common.py index ce03e7d333ddf..9be624d4cc1aa 100644 --- a/pandas/core/ops/common.py +++ b/pandas/core/ops/common.py @@ -46,11 +46,12 @@ 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=None): - if other is None: + if is_unary: return method(self) if is_cmp and isinstance(self, ABCIndexClass) and isinstance(other, ABCSeries): From df90be44d9bd7c723cd06b594fbb8415f73e0357 Mon Sep 17 00:00:00 2001 From: Asish Mahapatra Date: Wed, 2 Sep 2020 17:20:08 -0400 Subject: [PATCH 3/5] black --- pandas/core/arrays/integer.py | 2 +- pandas/core/ops/common.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pandas/core/arrays/integer.py b/pandas/core/arrays/integer.py index c81b6ae77141a..4568bc72892ab 100644 --- a/pandas/core/arrays/integer.py +++ b/pandas/core/arrays/integer.py @@ -653,7 +653,7 @@ 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__ diff --git a/pandas/core/ops/common.py b/pandas/core/ops/common.py index 9be624d4cc1aa..349e699426cf0 100644 --- a/pandas/core/ops/common.py +++ b/pandas/core/ops/common.py @@ -46,7 +46,7 @@ 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'} + is_unary = name.strip("__") in {"neg", "pos"} @wraps(method) def new_method(self, other=None): From 5c2a9e9a5c6ef607d5011209e1d59cf844a99b24 Mon Sep 17 00:00:00 2001 From: Asish Mahapatra Date: Wed, 2 Sep 2020 17:53:26 -0400 Subject: [PATCH 4/5] add tests --- pandas/tests/extension/base/ops.py | 12 ++++++++++++ pandas/tests/extension/test_integer.py | 4 ++++ 2 files changed, 16 insertions(+) diff --git a/pandas/tests/extension/base/ops.py b/pandas/tests/extension/base/ops.py index c93603398977e..7d87b64382a90 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_negate(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 From d53e5c21cf7445e38e2501c0e237e75b5b7c680e Mon Sep 17 00:00:00 2001 From: Asish Mahapatra Date: Wed, 2 Sep 2020 17:57:48 -0400 Subject: [PATCH 5/5] clean up names --- pandas/core/arrays/integer.py | 4 ++-- pandas/tests/extension/base/ops.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pandas/core/arrays/integer.py b/pandas/core/arrays/integer.py index 4568bc72892ab..8450dd50ba461 100644 --- a/pandas/core/arrays/integer.py +++ b/pandas/core/arrays/integer.py @@ -659,14 +659,14 @@ def _create_unary_method(cls, op): op_name = op.__name__ @unpack_zerodim_and_defer(op.__name__) - def unary_arithmetic_method(self): + 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(unary_arithmetic_method, name, cls) + return set_function_name(integer_unary_method, name, cls) IntegerArray._add_arithmetic_ops() diff --git a/pandas/tests/extension/base/ops.py b/pandas/tests/extension/base/ops.py index 7d87b64382a90..bbcace9e3e096 100644 --- a/pandas/tests/extension/base/ops.py +++ b/pandas/tests/extension/base/ops.py @@ -187,7 +187,7 @@ def test_invert(self, data): expected = pd.Series(~data, name="name") self.assert_series_equal(result, expected) - def test_negate(self, data): + def test_neg(self, data): s = pd.Series(data, name="name") result = -s expected = pd.Series(-data, name="name")