REF: use idiomatic checks in __array_ufunc__ (#44765)

Author: jbrockmendel

pandas/core/arraylike.py

@@ -326,13 +326,16 @@ def array_ufunc(self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: Any)
         reconstruct_kwargs = {}
 
     def reconstruct(result):
+        if ufunc.nout > 1:
+            # np.modf, np.frexp, np.divmod
+            return tuple(_reconstruct(x) for x in result)
+
+        return _reconstruct(result)
+
+    def _reconstruct(result):
         if lib.is_scalar(result):
             return result
 
-        if isinstance(result, tuple):
-            # np.modf, np.frexp, np.divmod
-            return tuple(reconstruct(x) for x in result)
-
         if result.ndim != self.ndim:
             if method == "outer":
                 if self.ndim == 2:

pandas/core/arrays/masked.py

@@ -479,7 +479,8 @@ def reconstruct(x):
             return x
 
         result = getattr(ufunc, method)(*inputs2, **kwargs)
-        if isinstance(result, tuple):
+        if ufunc.nout > 1:
+            # e.g. np.divmod
             return tuple(reconstruct(x) for x in result)
         else:
             return reconstruct(result)

pandas/core/arrays/numpy_.py

@@ -152,7 +152,7 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
             )
         result = getattr(ufunc, method)(*inputs, **kwargs)
 
-        if type(result) is tuple and len(result):
+        if ufunc.nout > 1:
             # multiple return values
             if not lib.is_scalar(result[0]):
                 # re-box array-like results
@@ -163,6 +163,13 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
         elif method == "at":
             # no return value
             return None
+        elif method == "reduce":
+            if isinstance(result, np.ndarray):
+                # e.g. test_np_reduce_2d
+                return type(self)(result)
+
+            # e.g. test_np_max_nested_tuples
+            return result
         else:
             # one return value
             if not lib.is_scalar(result):

pandas/core/arrays/sparse/array.py

@@ -1580,7 +1580,7 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
             sp_values = getattr(ufunc, method)(self.sp_values, **kwargs)
             fill_value = getattr(ufunc, method)(self.fill_value, **kwargs)
 
-            if isinstance(sp_values, tuple):
+            if ufunc.nout > 1:
                 # multiple outputs. e.g. modf
                 arrays = tuple(
                     self._simple_new(
@@ -1589,7 +1589,7 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
                     for sp_value, fv in zip(sp_values, fill_value)
                 )
                 return arrays
-            elif is_scalar(sp_values):
+            elif method == "reduce":
                 # e.g. reductions
                 return sp_values
 
@@ -1603,7 +1603,7 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
                 out = out[0]
             return out
 
-        if type(result) is tuple:
+        if ufunc.nout > 1:
             return tuple(type(self)(x) for x in result)
         elif method == "at":
             # no return value

pandas/tests/arrays/test_numpy.py

@@ -194,6 +194,38 @@ def test_validate_reduction_keyword_args():
         arr.all(keepdims=True)
 
 
+def test_np_max_nested_tuples():
+    # case where checking in ufunc.nout works while checking for tuples
+    #  does not
+    vals = [
+        (("j", "k"), ("l", "m")),
+        (("l", "m"), ("o", "p")),
+        (("o", "p"), ("j", "k")),
+    ]
+    ser = pd.Series(vals)
+    arr = ser.array
+
+    assert arr.max() is arr[2]
+    assert ser.max() is arr[2]
+
+    result = np.maximum.reduce(arr)
+    assert result == arr[2]
+
+    result = np.maximum.reduce(ser)
+    assert result == arr[2]
+
+
+def test_np_reduce_2d():
+    raw = np.arange(12).reshape(4, 3)
+    arr = PandasArray(raw)
+
+    res = np.maximum.reduce(arr, axis=0)
+    tm.assert_extension_array_equal(res, arr[-1])
+
+    alt = arr.max(axis=0)
+    tm.assert_extension_array_equal(alt, arr[-1])
+
+
 # ----------------------------------------------------------------------------
 # Ops
 

pandas/tests/extension/decimal/array.py

@@ -124,7 +124,7 @@ def reconstruct(x):
             else:
                 return DecimalArray._from_sequence(x)
 
-        if isinstance(result, tuple):
+        if ufunc.nout > 1:
             return tuple(reconstruct(x) for x in result)
         else:
             return reconstruct(result)