diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/src/private/ufunc_override.h | 13 | ||||
| -rw-r--r-- | numpy/core/tests/test_umath.py | 9 |
2 files changed, 19 insertions, 3 deletions
diff --git a/numpy/core/src/private/ufunc_override.h b/numpy/core/src/private/ufunc_override.h index d445ac2b8..380aef714 100644 --- a/numpy/core/src/private/ufunc_override.h +++ b/numpy/core/src/private/ufunc_override.h @@ -94,10 +94,17 @@ PyUFunc_CheckOverride(PyObject *ufunc, char *method, goto fail; } - /* If we have more args than nin, the last one must be `out`.*/ + /* If we have more args than nin, they must be the output variables.*/ if (nargs > nin) { - obj = PyTuple_GET_ITEM(args, nargs - 1); - PyDict_SetItemString(normal_kwds, "out", obj); + if ((nargs - nin) == 1) { + obj = PyTuple_GET_ITEM(args, nargs - 1); + PyDict_SetItemString(normal_kwds, "out", obj); + } + else { + obj = PyTuple_GetSlice(args, nin, nargs); + PyDict_SetItemString(normal_kwds, "out", obj); + Py_DECREF(obj); + } } method_name = PyUString_FromString(method); diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py index d61b516ac..6db95e9a6 100644 --- a/numpy/core/tests/test_umath.py +++ b/numpy/core/tests/test_umath.py @@ -1057,6 +1057,15 @@ class TestSpecialMethods(TestCase): assert_equal(res4['out'], 'out_arg') assert_equal(res5['out'], 'out_arg') + # ufuncs with multiple output modf and frexp. + res6 = np.modf(a, 'out0', 'out1') + res7 = np.frexp(a, 'out0', 'out1') + assert_equal(res6['out'][0], 'out0') + assert_equal(res6['out'][1], 'out1') + assert_equal(res7['out'][0], 'out0') + assert_equal(res7['out'][1], 'out1') + + def test_ufunc_override_exception(self): class A(object): def __numpy_ufunc__(self, *a, **kwargs): |