Merge pull request numpy#5830 from jaimefrio/frompyfunc_cleanup

MANT: cleanup logic in PyUFunc_On_Om
miclark · Aug 15, 2015 · c573b71 · c573b71
2 parents 59bbc06 + 13b5ce9
commit c573b71
Showing 1 changed file with 17 additions and 18 deletions.
diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src
@@ -603,34 +603,33 @@ PyUFunc_On_Om(char **args, npy_intp *dimensions, npy_intp *steps, void *func)
         if (result == NULL) {
             return;
         }
-        if (PyTuple_Check(result)) {
-            if (nout != PyTuple_Size(result)) {
-                Py_DECREF(result);
-                return;
-            }
-            for(j = 0; j < nout; j++) {
-                op = (PyObject **)ptrs[j+nin];
-                Py_XDECREF(*op);
-                *op = PyTuple_GET_ITEM(result, j);
-                Py_INCREF(*op);
-            }
+        if (nout == 0  && result == Py_None) {
+            /* No output expected, no output received, continue */
             Py_DECREF(result);
         }
         else if (nout == 1) {
+            /* Single output expected, assign and continue */
             op = (PyObject **)ptrs[nin];
             Py_XDECREF(*op);
             *op = result;
         }
-        else {
+        else if (PyTuple_Check(result) && nout == PyTuple_Size(result)) {
             /*
-             * The single output does not match the expected nout != 1,
-             * which may be reasonable for nout == 0. For other values,
-             * behave as if a 1-tuple had been returned and exit.
+             * Multiple returns match expected number of outputs, assign
+             * and continue. Will also gobble empty tuples if nout == 0.
              */
-            Py_DECREF(result);
-            if (nout != 0) {
-                return;
+            for(j = 0; j < nout; j++) {
+                op = (PyObject **)ptrs[j+nin];
+                Py_XDECREF(*op);
+                *op = PyTuple_GET_ITEM(result, j);
+                Py_INCREF(*op);
             }
+            Py_DECREF(result);
+        }
+        else {
+            /* Mismatch between returns and expected outputs, exit */
+            Py_DECREF(result);
+            return;
         }
         for(j = 0; j < ntot; j++) {
             ptrs[j] += steps[j];