Implement NA.__array_ufunc__

This gives us consistent comparisions with NumPy scalars.

Author: TomAugspurger

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

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,

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

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)