Add __array_ufunc__ to Series / Array

pandas/core/arrays/integer.py

@@ -1,3 +1,4 @@
+import numbers
 import sys
 import warnings
 import copy
@@ -293,6 +294,33 @@ def __array__(self, dtype=None):
         """
         return self._coerce_to_ndarray()
 
+    _HANDLED_TYPES = (np.ndarray, numbers.Number)
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+
+        out = kwargs.get('out', ())
+
+        for x in inputs + out:
+            if not isinstance(x, self._HANDLED_TYPES + (IntegerArray,)):
+                return NotImplemented
+
+        if method == '__call__':
+            if ufunc.signature is None and ufunc.nout == 1:
+                args = [a._data for a in inputs]
+                masks = [a._mask for a in inputs]
+                result = ufunc(*args, **kwargs)
+                mask = np.logical_or.reduce(masks)
+                if result.dtype.kind in ('i', 'u'):
+                    return IntegerArray(result, mask)
+                else:
+                    result[mask] = np.nan
+                    return result
+
+        # fall back to array for other ufuncs
+        return np.array(self).__array_ufunc__(
+            ufunc, method, *inputs, **kwargs)
+        return NotImplemented
+
     def __iter__(self):
         for i in range(len(self)):
             if self._mask[i]:

pandas/core/series.py

@@ -623,6 +623,23 @@ def view(self, dtype=None):
         return self._constructor(self._values.view(dtype),
                                  index=self.index).__finalize__(self)
 
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        inputs = tuple(
+            x._values if isinstance(x, type(self)) else x
+            for x in inputs
+        )
+        if hasattr(self._values, '__array_ufunc__'):
+            result = self._values.__array_ufunc__(
+                ufunc, method, *inputs, **kwargs)
+        else:
+            result = np.array(self._values).__array_ufunc__(
+                ufunc, method, *inputs, **kwargs)
+        if result is NotImplemented:
+            raise TypeError("The '{0}' operation is not supported for "
+                            "dtype {1}.".format(ufunc.__name__, self.dtype))
+        return self._constructor(result, index=self.index,
+                                 copy=False).__finalize__(self)
+
     def __array__(self, result=None):
         """
         the array interface, return my values
@@ -640,10 +657,9 @@ def __array_prepare__(self, result, context=None):
         """
         Gets called prior to a ufunc
         """
-
         # nice error message for non-ufunc types
         if (context is not None and
-                not isinstance(self._values, (np.ndarray, ABCSparseArray))):
+                not isinstance(self._values, (np.ndarray, ExtensionArray))):
             obj = context[1][0]
             raise TypeError("{obj} with dtype {dtype} cannot perform "
                             "the numpy op {op}".format(

pandas/tests/arrays/test_integer.py

@@ -690,6 +690,43 @@ def test_astype_nansafe():
         arr.astype('uint32')
 
 
+@pytest.mark.parametrize(
+    'ufunc', [np.abs, ])
+def test_ufuncs_single_int(ufunc):
+    a = integer_array([1, 2, -3, np.nan])
+    result = ufunc(a)
+    expected = integer_array(ufunc(a.astype(float)))
+    tm.assert_extension_array_equal(result, expected)
+
+    s = pd.Series(a)
+    result = ufunc(s)
+    expected = pd.Series(integer_array(ufunc(a.astype(float))))
+    tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    'ufunc', [np.log, np.exp, np.sin, np.cos, np.sqrt])
+def test_ufuncs_single_float(ufunc):
+    a = integer_array([1, 2, -3, np.nan])
+    result = ufunc(a)
+    expected = ufunc(a.astype(float))
+    tm.assert_numpy_array_equal(result, expected)
+
+    s = pd.Series(a)
+    result = ufunc(s)
+    expected = ufunc(s.astype(float))
+    tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    'ufunc', [np.add, np.subtract])
+def test_ufuncs_binary_int(ufunc):
+    a = integer_array([1, 2, -3, np.nan])
+    result = ufunc(a, a)
+    expected = integer_array(ufunc(a.astype(float), a.astype(float)))
+    tm.assert_extension_array_equal(result, expected)
+
+
 # TODO(jreback) - these need testing / are broken
 
 # shift