diff --git a/pandas/_libs/missing.pyx b/pandas/_libs/missing.pyx index 63aa5501c5250..e2503893e705d 100644 --- a/pandas/_libs/missing.pyx +++ b/pandas/_libs/missing.pyx @@ -290,7 +290,8 @@ def _create_binary_propagating_op(name, divmod=False): def method(self, other): if (other is C_NA or isinstance(other, str) - or isinstance(other, (numbers.Number, np.bool_))): + or isinstance(other, (numbers.Number, np.bool_, np.int64, np.int_)) + or isinstance(other, np.ndarray) and not other.shape): if divmod: return NA, NA else: @@ -310,6 +311,98 @@ def _create_unary_propagating_op(name): return method +def maybe_dispatch_ufunc_to_dunder_op( + object self, object ufunc, str method, *inputs, **kwargs +): + """ + Dispatch a ufunc to the equivalent dunder method. + + Parameters + ---------- + self : ArrayLike + The array whose dunder method we dispatch to + ufunc : Callable + A NumPy ufunc + method : {'reduce', 'accumulate', 'reduceat', 'outer', 'at', '__call__'} + inputs : ArrayLike + The input arrays. + kwargs : Any + The additional keyword arguments, e.g. ``out``. + + Returns + ------- + result : Any + The result of applying the ufunc + """ + # special has the ufuncs we dispatch to the dunder op on + special = { + "add", + "sub", + "mul", + "pow", + "mod", + "floordiv", + "truediv", + "divmod", + "eq", + "ne", + "lt", + "gt", + "le", + "ge", + "remainder", + "matmul", + "or", + "xor", + "and", + } + aliases = { + "subtract": "sub", + "multiply": "mul", + "floor_divide": "floordiv", + "true_divide": "truediv", + "power": "pow", + "remainder": "mod", + "divide": "div", + "equal": "eq", + "not_equal": "ne", + "less": "lt", + "less_equal": "le", + "greater": "gt", + "greater_equal": "ge", + "bitwise_or": "or", + "bitwise_and": "and", + "bitwise_xor": "xor", + } + + # For op(., Array) -> Array.__r{op}__ + flipped = { + "lt": "__gt__", + "le": "__ge__", + "gt": "__lt__", + "ge": "__le__", + "eq": "__eq__", + "ne": "__ne__", + } + + op_name = ufunc.__name__ + op_name = aliases.get(op_name, op_name) + + def not_implemented(*args, **kwargs): + return NotImplemented + + if method == "__call__" and op_name in special and kwargs.get("out") is None: + if isinstance(inputs[0], type(self)): + name = "__{}__".format(op_name) + return getattr(self, name, not_implemented)(inputs[1]) + else: + name = flipped.get(op_name, "__r{}__".format(op_name)) + result = getattr(self, name, not_implemented)(inputs[0]) + return result + else: + return NotImplemented + + cdef class C_NAType: pass @@ -434,6 +527,27 @@ class NAType(C_NAType): __rxor__ = __xor__ + # What else to add here? datetime / Timestamp? Period? Interval? + # Note: we only handle 0-d ndarrays. + __array_priority__ = 1000 + _HANDLED_TYPES = (np.ndarray, numbers.Number, str, np.bool_, np.int64) + + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + types = self._HANDLED_TYPES + (NAType,) + for x in inputs: + if not isinstance(x, types): + return NotImplemented + + if method != "__call__": + raise ValueError(f"ufunc method '{method}' not supported for NA") + result = maybe_dispatch_ufunc_to_dunder_op(self, ufunc, method, *inputs, **kwargs) + if result is NotImplemented: + if ufunc.nout == 1: + result = NA + else: + result = (NA,) * ufunc.nout + return result + C_NA = NAType() # C-visible NA = C_NA # Python-visible diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index f3c01efed6d43..eb097831e04d4 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -10,6 +10,7 @@ import numpy as np from pandas._libs import Timedelta, Timestamp, lib +from pandas._libs.missing import maybe_dispatch_ufunc_to_dunder_op # noqa:F401 from pandas.util._decorators import Appender from pandas.core.dtypes.common import is_list_like, is_timedelta64_dtype @@ -30,7 +31,6 @@ ) from pandas.core.ops.array_ops import comp_method_OBJECT_ARRAY # noqa:F401 from pandas.core.ops.common import unpack_zerodim_and_defer -from pandas.core.ops.dispatch import maybe_dispatch_ufunc_to_dunder_op # noqa:F401 from pandas.core.ops.dispatch import should_series_dispatch from pandas.core.ops.docstrings import ( _arith_doc_FRAME, diff --git a/pandas/core/ops/dispatch.py b/pandas/core/ops/dispatch.py index 016a89eb56da3..0c27d09a0b12d 100644 --- a/pandas/core/ops/dispatch.py +++ b/pandas/core/ops/dispatch.py @@ -1,7 +1,7 @@ """ Functions for defining unary operations. """ -from typing import Any, Callable, Union +from typing import Any, Union import numpy as np @@ -17,7 +17,6 @@ ) from pandas.core.dtypes.generic import ABCExtensionArray, ABCSeries -from pandas._typing import ArrayLike from pandas.core.construction import array @@ -146,94 +145,3 @@ def dispatch_to_extension_op( "operation [{name}]".format(name=op.__name__) ) return res_values - - -def maybe_dispatch_ufunc_to_dunder_op( - self: ArrayLike, ufunc: Callable, method: str, *inputs: ArrayLike, **kwargs: Any -): - """ - Dispatch a ufunc to the equivalent dunder method. - - Parameters - ---------- - self : ArrayLike - The array whose dunder method we dispatch to - ufunc : Callable - A NumPy ufunc - method : {'reduce', 'accumulate', 'reduceat', 'outer', 'at', '__call__'} - inputs : ArrayLike - The input arrays. - kwargs : Any - The additional keyword arguments, e.g. ``out``. - - Returns - ------- - result : Any - The result of applying the ufunc - """ - # special has the ufuncs we dispatch to the dunder op on - special = { - "add", - "sub", - "mul", - "pow", - "mod", - "floordiv", - "truediv", - "divmod", - "eq", - "ne", - "lt", - "gt", - "le", - "ge", - "remainder", - "matmul", - "or", - "xor", - "and", - } - aliases = { - "subtract": "sub", - "multiply": "mul", - "floor_divide": "floordiv", - "true_divide": "truediv", - "power": "pow", - "remainder": "mod", - "divide": "div", - "equal": "eq", - "not_equal": "ne", - "less": "lt", - "less_equal": "le", - "greater": "gt", - "greater_equal": "ge", - "bitwise_or": "or", - "bitwise_and": "and", - "bitwise_xor": "xor", - } - - # For op(., Array) -> Array.__r{op}__ - flipped = { - "lt": "__gt__", - "le": "__ge__", - "gt": "__lt__", - "ge": "__le__", - "eq": "__eq__", - "ne": "__ne__", - } - - op_name = ufunc.__name__ - op_name = aliases.get(op_name, op_name) - - def not_implemented(*args, **kwargs): - return NotImplemented - - if method == "__call__" and op_name in special and kwargs.get("out") is None: - if isinstance(inputs[0], type(self)): - name = "__{}__".format(op_name) - return getattr(self, name, not_implemented)(inputs[1]) - else: - name = flipped.get(op_name, "__r{}__".format(op_name)) - return getattr(self, name, not_implemented)(inputs[0]) - else: - return NotImplemented diff --git a/pandas/tests/scalar/test_na_scalar.py b/pandas/tests/scalar/test_na_scalar.py index 40db617c64717..a65cb6a0cbb01 100644 --- a/pandas/tests/scalar/test_na_scalar.py +++ b/pandas/tests/scalar/test_na_scalar.py @@ -58,12 +58,6 @@ def test_comparison_ops(): assert (NA >= other) is NA assert (NA < other) is NA assert (NA <= other) is NA - - if isinstance(other, (np.int64, np.bool_)): - # for numpy scalars we get a deprecation warning and False as result - # for equality or error for larger/lesser than - continue - assert (other == NA) is NA assert (other != NA) is NA assert (other > NA) is NA @@ -175,3 +169,18 @@ def test_series_isna(): s = pd.Series([1, NA], dtype=object) expected = pd.Series([False, True]) tm.assert_series_equal(s.isna(), expected) + + +def test_ufunc(): + assert np.log(pd.NA) is pd.NA + assert np.add(pd.NA, 1) is pd.NA + result = np.divmod(pd.NA, 1) + assert result[0] is pd.NA and result[1] is pd.NA + + result = np.frexp(pd.NA) + assert result[0] is pd.NA and result[1] is pd.NA + + +def test_ufunc_raises(): + with pytest.raises(ValueError, match="ufunc method 'at'"): + np.log.at(pd.NA, 0)