diff --git a/numpy/core/src/private/ufunc_override.h b/numpy/core/src/private/ufunc_override.h index c3f9f601e0ff..4042eae2fde2 100644 --- 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 index ac645f01322c..0f024cbf721d 100644 --- 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):