ENH: ufuncs can take a tuple of arrays as 'out' kwarg

Closes #4752

4 files changed, 311 insertions, 100 deletions

diff --git a/doc/release/1.10.0-notes.rst b/doc/release/1.10.0-notes.rst
index a07dca80f..fe053af5f 100644
--- a/doc/release/1.10.0-notes.rst
+++ b/doc/release/1.10.0-notes.rst
@@ -71,6 +71,15 @@ Notably, this affect recarrays containing strings with whitespace, as trailing
 whitespace is trimmed from chararrays but kept in ndarrays of string type.
 Also, the dtype.type of nested structured fields is now inherited.
 
+'out' keyword argument of ufuncs now accepts tuples of arrays
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+When using the 'out' keyword argument of a ufunc, a tuple of arrays, one per
+ufunc output, can be provided. For ufuncs with a single output a single array
+is also a valid 'out' keyword argument. Previously a single array could be
+provided in the 'out' keyword argument, and it would be used as the first
+output for ufuncs with multiple outputs, is deprecated, and will result in a
+`DeprecationWarning` now and an error in the future.
+
 New Features
 ============
 
diff --git a/doc/source/reference/ufuncs.rst b/doc/source/reference/ufuncs.rst
index 3d6112058..f7142bc47 100644
--- a/doc/source/reference/ufuncs.rst
+++ b/doc/source/reference/ufuncs.rst
@@ -299,7 +299,20 @@ advanced usage and will not typically be used.
 
     .. versionadded:: 1.6
 
-    The first output can provided as either a positional or a keyword parameter.
+    The first output can be provided as either a positional or a keyword
+    parameter. Keyword 'out' arguments are incompatible with positional
+    ones.
+
+    ..versionadded:: 1.10
+
+    The 'out' keyword argument is expected to be a tuple with one entry per
+    output (which can be `None` for arrays to be allocated by the ufunc).
+    For ufuncs with a single output, passing a single array (instead of a
+    tuple holding a single array) is also valid.
+
+    Passing a single array in the 'out' keyword argument to a ufunc with
+    multiple outputs is deprecated, and will raise a warning in numpy 1.10,
+    and an error in a future release.
 
 *where*
 
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index dc5065f14..5e4a49553 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -750,6 +750,35 @@ fail:
     return -1;
 }
 
+/*
+ * Checks if 'obj' is a valid output array for a ufunc, i.e. it is
+ * either None or a writeable array, increments its reference count
+ * and stores a pointer to it in 'store'. Returns 0 on success, sets
+ * an exception and returns -1 on failure.
+ */
+static int
+_set_out_array(PyObject *obj, PyArrayObject **store)
+{
+    if (obj == Py_None) {
+        /* Translate None to NULL */
+        return 0;
+    }
+    if PyArray_Check(obj) {
+        /* If it's an array, store it */
+        if (PyArray_FailUnlessWriteable((PyArrayObject *)obj,
+                                        "output array") < 0) {
+            return -1;
+        }
+        Py_INCREF(obj);
+        *store = (PyArrayObject *)obj;
+
+        return 0;
+    }
+    PyErr_SetString(PyExc_TypeError, "return arrays must be of ArrayType");
+
+    return -1;
+}
+
 /********* GENERIC UFUNC USING ITERATOR *********/
 
 /*
@@ -759,17 +788,20 @@ fail:
  * non-zero references in out_op.  This
  * function does not do its own clean-up.
  */
-static int get_ufunc_arguments(PyUFuncObject *ufunc,
-                PyObject *args, PyObject *kwds,
-                PyArrayObject **out_op,
-                NPY_ORDER *out_order,
-                NPY_CASTING *out_casting,
-                PyObject **out_extobj,
-                PyObject **out_typetup,
-                int *out_subok,
-                PyArrayObject **out_wheremask)
+static int
+get_ufunc_arguments(PyUFuncObject *ufunc,
+                    PyObject *args, PyObject *kwds,
+                    PyArrayObject **out_op,
+                    NPY_ORDER *out_order,
+                    NPY_CASTING *out_casting,
+                    PyObject **out_extobj,
+                    PyObject **out_typetup,
+                    int *out_subok,
+                    PyArrayObject **out_wheremask)
 {
-    int i, nargs, nin = ufunc->nin;
+    int i, nargs;
+    int nin = ufunc->nin;
+    int nout = ufunc->nout;
     PyObject *obj, *context;
     PyObject *str_key_obj = NULL;
     const char *ufunc_name;
@@ -878,23 +910,7 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
     /* Get positional output arguments */
     for (i = nin; i < nargs; ++i) {
         obj = PyTuple_GET_ITEM(args, i);
-        /* Translate None to NULL */
-        if (obj == Py_None) {
-            continue;
-        }
-        /* If it's an array, can use it */
-        if (PyArray_Check(obj)) {
-            if (PyArray_FailUnlessWriteable((PyArrayObject *)obj,
-                                            "output array") < 0) {
-                return -1;
-            }
-            Py_INCREF(obj);
-            out_op[i] = (PyArrayObject *)obj;
-        }
-        else {
-            PyErr_SetString(PyExc_TypeError,
-                            "return arrays must be "
-                            "of ArrayType");
+        if (_set_out_array(obj, out_op + i) < 0) {
             return -1;
         }
     }
@@ -929,7 +945,7 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
             switch (str[0]) {
                 case 'c':
                     /* Provides a policy for allowed casting */
-                    if (strncmp(str,"casting",7) == 0) {
+                    if (strncmp(str, "casting", 7) == 0) {
                         if (!PyArray_CastingConverter(value, out_casting)) {
                             goto fail;
                         }
@@ -938,7 +954,7 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
                     break;
                 case 'd':
                     /* Another way to specify 'sig' */
-                    if (strncmp(str,"dtype",5) == 0) {
+                    if (strncmp(str, "dtype", 5) == 0) {
                         /* Allow this parameter to be None */
                         PyArray_Descr *dtype;
                         if (!PyArray_DescrConverter2(value, &dtype)) {
@@ -960,35 +976,74 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
                      * Overrides the global parameters buffer size,
                      * error mask, and error object
                      */
-                    if (strncmp(str,"extobj",6) == 0) {
+                    if (strncmp(str, "extobj", 6) == 0) {
                         *out_extobj = value;
                         bad_arg = 0;
                     }
                     break;
                 case 'o':
-                    /* First output may be specified as a keyword parameter */
-                    if (strncmp(str,"out",3) == 0) {
-                        if (out_op[nin] != NULL) {
+                    /*
+                     * Output arrays may be specified as a keyword argument,
+                     * either as a single array or None for single output
+                     * ufuncs, or as a tuple of arrays and Nones.
+                     */
+                    if (strncmp(str, "out", 3) == 0) {
+                        if (nargs > nin) {
                             PyErr_SetString(PyExc_ValueError,
                                     "cannot specify 'out' as both a "
                                     "positional and keyword argument");
                             goto fail;
                         }
-
-                        if (PyArray_Check(value)) {
-                            const char *name = "output array";
-                            PyArrayObject *value_arr = (PyArrayObject *)value;
-                            if (PyArray_FailUnlessWriteable(value_arr, name) < 0) {
+                        if (PyTuple_Check(value)) {
+                            if (PyTuple_GET_SIZE(value) != nout) {
+                                PyErr_SetString(PyExc_ValueError,
+                                        "The 'out' tuple must have exactly "
+                                        "one entry per ufunc output");
+                                goto fail;
+                            }
+                            /* 'out' must be a tuple of arrays and Nones */
+                            for(i = 0; i < nout; ++i) {
+                                PyObject *val = PyTuple_GET_ITEM(value, i);
+                                if (_set_out_array(val, out_op+nin+i) < 0) {
+                                    goto fail;
+                                }
+                            }
+                        }
+                        else if (nout == 1) {
+                            /* Can be an array if it only has one output */
+                            if (_set_out_array(value, out_op + nin) < 0) {
                                 goto fail;
                             }
-                            Py_INCREF(value);
-                            out_op[nin] = (PyArrayObject *)value;
                         }
                         else {
-                            PyErr_SetString(PyExc_TypeError,
-                                            "return arrays must be "
-                                            "of ArrayType");
-                            goto fail;
+                            /*
+                             * If the deprecated behavior is ever removed,
+                             * keep only the else branch of this if-else
+                             */
+                            if (PyArray_Check(value) || value == Py_None) {
+                                if (DEPRECATE("passing a single array to the "
+                                              "'out' keyword argument of a "
+                                              "ufunc with\n"
+                                              "more than one output will "
+                                              "result in an error in the "
+                                              "future") < 0) {
+                                    /* The future error message */
+                                    PyErr_SetString(PyExc_TypeError,
+                                        "'out' must be a tuple of arrays");
+                                    goto fail;
+                                }
+                                if (_set_out_array(value, out_op+nin) < 0) {
+                                    goto fail;
+                                }
+                            }
+                            else {
+                                PyErr_SetString(PyExc_TypeError,
+                                    nout > 1 ? "'out' must be a tuple "
+                                               "of arrays" :
+                                               "'out' must be an array or a "
+                                               "tuple of a single array");
+                                goto fail;
+                            }
                         }
                         bad_arg = 0;
                     }
@@ -3946,6 +4001,38 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args,
 }
 
 /*
+ * Returns an incref'ed pointer to the proper wrapping object for a
+ * ufunc output argument, given the output argument 'out', and the
+ * input's wrapping function, 'wrap'.
+ */
+static PyObject*
+_get_out_wrap(PyObject *out, PyObject *wrap) {
+    PyObject *owrap;
+
+    if (out == Py_None) {
+        /* Iterator allocated outputs get the input's wrapping */
+        Py_XINCREF(wrap);
+        return wrap;
+    }
+    if (PyArray_CheckExact(out)) {
+        /* None signals to not call any wrapping */
+        Py_RETURN_NONE;
+    }
+    /*
+     * For array subclasses use their __array_wrap__ method, or the
+     * input's wrapping if not available
+     */
+    owrap = PyObject_GetAttr(out, npy_um_str_array_wrap);
+    if (owrap == NULL || !PyCallable_Check(owrap)) {
+        Py_XDECREF(owrap);
+        owrap = wrap;
+        Py_XINCREF(wrap);
+        PyErr_Clear();
+    }
+    return owrap;
+}
+
+/*
  * This function analyzes the input arguments
  * and determines an appropriate __array_wrap__ function to call
  * for the outputs.
@@ -3966,7 +4053,7 @@ _find_array_wrap(PyObject *args, PyObject *kwds,
                 PyObject **output_wrap, int nin, int nout)
 {
     Py_ssize_t nargs;
-    int i;
+    int i, idx_offset, start_idx;
     int np = 0;
     PyObject *with_wrap[NPY_MAXARGS], *wraps[NPY_MAXARGS];
     PyObject *obj, *wrap = NULL;
@@ -4043,45 +4130,45 @@ _find_array_wrap(PyObject *args, PyObject *kwds,
      */
 handle_out:
     nargs = PyTuple_GET_SIZE(args);
-    for (i = 0; i < nout; i++) {
-        int j = nin + i;
-        int incref = 1;
-        output_wrap[i] = wrap;
-        obj = NULL;
-        if (j < nargs) {
-            obj = PyTuple_GET_ITEM(args, j);
-            /* Output argument one may also be in a keyword argument */
-            if (i == 0 && obj == Py_None && kwds != NULL) {
-                obj = PyDict_GetItem(kwds, npy_um_str_out);
-            }
+    /* Default is using positional arguments */
+    obj = args;
+    idx_offset = nin;
+    start_idx = 0;
+    if (nin == nargs && kwds != NULL) {
+        /* There may be a keyword argument we can use instead */
+        obj = PyDict_GetItem(kwds, npy_um_str_out);
+        if (obj == NULL) {
+            /* No, go back to positional (even though there aren't any) */
+            obj = args;
         }
-        /* Output argument one may also be in a keyword argument */
-        else if (i == 0 && kwds != NULL) {
-            obj = PyDict_GetItem(kwds, npy_um_str_out);
-        }
-
-        if (obj != Py_None && obj != NULL) {
-            if (PyArray_CheckExact(obj)) {
-                /* None signals to not call any wrapping */
-                output_wrap[i] = Py_None;
+        else {
+            idx_offset = 0;
+            if (PyTuple_Check(obj)) {
+                /* If a tuple, must have all nout items */
+                nargs = nout;
             }
             else {
-                PyObject *owrap = PyObject_GetAttr(obj, npy_um_str_array_wrap);
-                incref = 0;
-                if (!(owrap) || !(PyCallable_Check(owrap))) {
-                    Py_XDECREF(owrap);
-                    owrap = wrap;
-                    incref = 1;
-                    PyErr_Clear();
-                }
-                output_wrap[i] = owrap;
+                /* If the kwarg is not a tuple then it is an array (or None) */
+                output_wrap[0] = _get_out_wrap(obj, wrap);
+                start_idx = 1;
+                nargs = 1;
             }
         }
+    }
 
-        if (incref) {
-            Py_XINCREF(output_wrap[i]);
+    for (i = start_idx; i < nout; ++i) {
+        int j = idx_offset + i;
+
+        if (j < nargs) {
+            output_wrap[i] = _get_out_wrap(PyTuple_GET_ITEM(obj, j),
+                                           wrap);
+        }
+        else {
+            output_wrap[i] = wrap;
+            Py_XINCREF(wrap);
         }
     }
+
     Py_XDECREF(wrap);
     return;
 }
diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
index 092953872..e8eee8090 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -36,35 +36,137 @@ class TestConstants(TestCase):
     def test_euler_gamma(self):
         assert_allclose(ncu.euler_gamma, 0.5772156649015329, 1e-15)
 
+
 class TestOut(TestCase):
     def test_out_subok(self):
-        for b in (True, False):
-            aout = np.array(0.5)
-
-            r = np.add(aout, 2, out=aout)
-            assert_(r is aout)
-            assert_array_equal(r, aout)
-
-            r = np.add(aout, 2, out=aout, subok=b)
-            assert_(r is aout)
-            assert_array_equal(r, aout)
-
-            r = np.add(aout, 2, aout, subok=False)
-            assert_(r is aout)
-            assert_array_equal(r, aout)
-
-            d = np.ones(5)
-            o1 = np.zeros(5)
-            o2 = np.zeros(5, dtype=np.int32)
-            r1, r2 = np.frexp(d, o1, o2, subok=b)
+        for subok in (True, False):
+            a = np.array(0.5)
+            o = np.empty(())
+
+            r = np.add(a, 2, o, subok=subok)
+            assert_(r is o)
+            r = np.add(a, 2, out=o, subok=subok)
+            assert_(r is o)
+            r = np.add(a, 2, out=(o,), subok=subok)
+            assert_(r is o)
+
+            d = np.array(5.7)
+            o1 = np.empty(())
+            o2 = np.empty((), dtype=np.int32)
+
+            r1, r2 = np.frexp(d, o1, None, subok=subok)
+            assert_(r1 is o1)
+            r1, r2 = np.frexp(d, None, o2, subok=subok)
+            assert_(r2 is o2)
+            r1, r2 = np.frexp(d, o1, o2, subok=subok)
             assert_(r1 is o1)
-            assert_array_equal(r1, o1)
             assert_(r2 is o2)
-            assert_array_equal(r2, o2)
 
-            r1, r2 = np.frexp(d, out=o1, subok=b)
+            r1, r2 = np.frexp(d, out=(o1, None), subok=subok)
             assert_(r1 is o1)
-            assert_array_equal(r1, o1)
+            r1, r2 = np.frexp(d, out=(None, o2), subok=subok)
+            assert_(r2 is o2)
+            r1, r2 = np.frexp(d, out=(o1, o2), subok=subok)
+            assert_(r1 is o1)
+            assert_(r2 is o2)
+
+            with warnings.catch_warnings(record=True) as w:
+                warnings.filterwarnings('always', '', DeprecationWarning)
+                r1, r2 = np.frexp(d, out=o1, subok=subok)
+                assert_(r1 is o1)
+                assert_(w[0].category is DeprecationWarning)
+
+            assert_raises(ValueError, np.add, a, 2, o, o, subok=subok)
+            assert_raises(ValueError, np.add, a, 2, o, out=o, subok=subok)
+            assert_raises(ValueError, np.add, a, 2, None, out=o, subok=subok)
+            assert_raises(ValueError, np.add, a, 2, out=(o, o), subok=subok)
+            assert_raises(ValueError, np.add, a, 2, out=(), subok=subok)
+            assert_raises(TypeError, np.add, a, 2, [], subok=subok)
+            assert_raises(TypeError, np.add, a, 2, out=[], subok=subok)
+            assert_raises(TypeError, np.add, a, 2, out=([],), subok=subok)
+            o.flags.writeable = False
+            assert_raises(ValueError, np.add, a, 2, o, subok=subok)
+            assert_raises(ValueError, np.add, a, 2, out=o, subok=subok)
+            assert_raises(ValueError, np.add, a, 2, out=(o,), subok=subok)
+
+
+    def test_out_wrap_subok(self):
+        class ArrayWrap(np.ndarray):
+            __array_priority__ = 10
+            def __new__(cls, arr):
+                return np.asarray(arr).view(cls).copy()
+            def __array_wrap__(self, arr, context):
+                return arr.view(type(self))
+
+        for subok in (True, False):
+            a = ArrayWrap([0.5])
+
+            r = np.add(a, 2, subok=subok)
+            if subok:
+                assert_(isinstance(r, ArrayWrap))
+            else:
+                assert_(type(r) == np.ndarray)
+
+            r = np.add(a, 2, None, subok=subok)
+            if subok:
+                assert_(isinstance(r, ArrayWrap))
+            else:
+                assert_(type(r) == np.ndarray)
+
+            r = np.add(a, 2, out=None, subok=subok)
+            if subok:
+                assert_(isinstance(r, ArrayWrap))
+            else:
+                assert_(type(r) == np.ndarray)
+
+            r = np.add(a, 2, out=(None,), subok=subok)
+            if subok:
+                assert_(isinstance(r, ArrayWrap))
+            else:
+                assert_(type(r) == np.ndarray)
+
+            d = ArrayWrap([5.7])
+            o1 = np.empty((1,))
+            o2 = np.empty((1,), dtype=np.int32)
+
+            r1, r2 = np.frexp(d, o1, subok=subok)
+            if subok:
+                assert_(isinstance(r2, ArrayWrap))
+            else:
+                assert_(type(r2) == np.ndarray)
+
+            r1, r2 = np.frexp(d, o1, None, subok=subok)
+            if subok:
+                assert_(isinstance(r2, ArrayWrap))
+            else:
+                assert_(type(r2) == np.ndarray)
+
+            r1, r2 = np.frexp(d, None, o2, subok=subok)
+            if subok:
+                assert_(isinstance(r1, ArrayWrap))
+            else:
+                assert_(type(r1) == np.ndarray)
+
+            r1, r2 = np.frexp(d, out=(o1, None), subok=subok)
+            if subok:
+                assert_(isinstance(r2, ArrayWrap))
+            else:
+                assert_(type(r2) == np.ndarray)
+
+            r1, r2 = np.frexp(d, out=(None, o2), subok=subok)
+            if subok:
+                assert_(isinstance(r1, ArrayWrap))
+            else:
+                assert_(type(r1) == np.ndarray)
+
+            with warnings.catch_warnings(record=True) as w:
+                warnings.filterwarnings('always', '', DeprecationWarning)
+                r1, r2 = np.frexp(d, out=o1, subok=subok)
+                if subok:
+                    assert_(isinstance(r2, ArrayWrap))
+                else:
+                    assert_(type(r2) == np.ndarray)
+                assert_(w[0].category is DeprecationWarning)
 
 
 class TestDivision(TestCase):