Check out kwarg for __nump_ufunc__ override and set index appropriately

cowlicks · mhvk · commit 288239d3485e · 2015-06-19T13:35:15.000-04:00
for the case where self is among outputs but not among inputs. Ensure it
works both out passed on as an argument and with out in a keyword argument.
diff --git a/numpy/core/src/private/ufunc_override.h b/numpy/core/src/private/ufunc_override.h
@@ -180,10 +180,13 @@ PyUFunc_CheckOverride(PyUFuncObject *ufunc, char *method,
     int override_pos; /* Position of override in args.*/
     int j;
 
-    int nargs = PyTuple_GET_SIZE(args);
+    int nargs;
+    int nout_kwd = 0;
+    int out_kwd_is_tuple = 0;
     int noa = 0; /* Number of overriding args.*/
 
     PyObject *obj;
+    PyObject *out_kwd_obj = NULL;
     PyObject *other_obj;
 
     PyObject *method_name = NULL;
@@ -204,16 +207,40 @@ PyUFunc_CheckOverride(PyUFuncObject *ufunc, char *method,
                         "with non-tuple");
         goto fail;
     }
-
-    if (PyTuple_GET_SIZE(args) > NPY_MAXARGS) {
+    nargs = PyTuple_GET_SIZE(args);
+    if (nargs > NPY_MAXARGS) {
         PyErr_SetString(PyExc_ValueError,
                         "Internal Numpy error: too many arguments in call "
                         "to PyUFunc_CheckOverride");
         goto fail;
     }
 
-    for (i = 0; i < nargs; ++i) {
-        obj = PyTuple_GET_ITEM(args, i);
+    /* be sure to include possible 'out' keyword argument. */
+    if ((kwds)&& (PyDict_CheckExact(kwds))) {
+        out_kwd_obj = PyDict_GetItemString(kwds, "out");
+        if (out_kwd_obj != NULL) {
+            out_kwd_is_tuple = PyTuple_CheckExact(out_kwd_obj);
+            if (out_kwd_is_tuple) {
+                nout_kwd = PyTuple_GET_SIZE(out_kwd_obj);
+            }
+            else {
+                nout_kwd = 1;
+            }
+        }
+    }
+
+    for (i = 0; i < nargs + nout_kwd; ++i) {
+        if (i < nargs) {
+            obj = PyTuple_GET_ITEM(args, i);
+        }
+        else {
+            if (out_kwd_is_tuple) {
+                obj = PyTuple_GET_ITEM(out_kwd_obj, i-nargs);
+            }
+            else {
+                obj = out_kwd_obj;
+            }
+        }
         /*
          * TODO: could use PyArray_GetAttrString_SuppressException if it
          * weren't private to multiarray.so
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
@@ -2383,15 +2383,15 @@ def __numpy_ufunc__(self, ufunc, method, i, inputs, **kw):
                     return "ufunc"
                 else:
                     inputs = list(inputs)
-                    inputs[i] = np.asarray(self)
+                    if i < len(inputs):
+                        inputs[i] = np.asarray(self)
                     func = getattr(ufunc, method)
+                    if ('out' in kw) and (kw['out'] is not None):
+                        kw['out'] = np.asarray(kw['out'])
                     r = func(*inputs, **kw)
-                    if 'out' in kw:
-                        return r
-                    else:
-                        x = self.__class__(r.shape, dtype=r.dtype)
-                        x[...] = r
-                        return x
+                    x = self.__class__(r.shape, dtype=r.dtype)
+                    x[...] = r
+                    return x
 
         class SomeClass3(SomeClass2):
             def __rsub__(self, other):
@@ -2475,6 +2475,64 @@ def __numpy_ufunc__(self, ufunc, method, i, inputs, **kw):
         assert_('sig' not in kw and 'signature' in kw)
         assert_equal(kw['signature'], 'ii->i')
 
+    def test_numpy_ufunc_index(self):
+        # Check that index is set appropriately, also if only an output
+        # is passed on (latter is another regression tests for github bug 4753)
+        class CheckIndex(object):
+            def __numpy_ufunc__(self, ufunc, method, i, inputs, **kw):
+                return i
+
+        a = CheckIndex()
+        dummy = np.arange(2.)
+        # 1 input, 1 output
+        assert_equal(np.sin(a), 0)
+        assert_equal(np.sin(dummy, a), 1)
+        assert_equal(np.sin(dummy, out=a), 1)
+        assert_equal(np.sin(dummy, out=(a,)), 1)
+        assert_equal(np.sin(a, a), 0)
+        assert_equal(np.sin(a, out=a), 0)
+        assert_equal(np.sin(a, out=(a,)), 0)
+        # 1 input, 2 outputs
+        assert_equal(np.modf(dummy, a), 1)
+        assert_equal(np.modf(dummy, None, a), 2)
+        assert_equal(np.modf(dummy, dummy, a), 2)
+        assert_equal(np.modf(dummy, out=a), 1)
+        assert_equal(np.modf(dummy, out=(a,)), 1)
+        assert_equal(np.modf(dummy, out=(a, None)), 1)
+        assert_equal(np.modf(dummy, out=(a, dummy)), 1)
+        assert_equal(np.modf(dummy, out=(None, a)), 2)
+        assert_equal(np.modf(dummy, out=(dummy, a)), 2)
+        assert_equal(np.modf(a, out=(dummy, a)), 0)
+        # 2 inputs, 1 output
+        assert_equal(np.add(a, dummy), 0)
+        assert_equal(np.add(dummy, a), 1)
+        assert_equal(np.add(dummy, dummy, a), 2)
+        assert_equal(np.add(dummy, a, a), 1)
+        assert_equal(np.add(dummy, dummy, out=a), 2)
+        assert_equal(np.add(dummy, dummy, out=(a,)), 2)
+        assert_equal(np.add(a, dummy, out=a), 0)
+
+    def test_out_override(self):
+        # regression test for github bug 4753
+        class OutClass(ndarray):
+            def __numpy_ufunc__(self, ufunc, method, i, inputs, **kw):
+                if 'out' in kw:
+                    tmp_kw = kw.copy()
+                    tmp_kw.pop('out')
+                    func = getattr(ufunc, method)
+                    kw['out'][...] = func(*inputs, **tmp_kw)
+
+        A = np.array([0]).view(OutClass)
+        B = np.array([5])
+        C = np.array([6])
+        np.multiply(C, B, A)
+        assert_equal(A[0], 30)
+        assert_(isinstance(A, OutClass))
+        A[0] = 0
+        np.multiply(C, B, out=A)
+        assert_equal(A[0], 30)
+        assert_(isinstance(A, OutClass))
+
 
 class TestCAPI(TestCase):
     def test_IsPythonScalar(self):
