diff options
Diffstat (limited to 'numpy/core')
| -rw-r--r-- | numpy/core/src/multiarray/nditer_pywrap.c | 2 | ||||
| -rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 2 | ||||
| -rw-r--r-- | numpy/core/tests/test_nditer.py | 23 |
3 files changed, 26 insertions, 1 deletions
diff --git a/numpy/core/src/multiarray/nditer_pywrap.c b/numpy/core/src/multiarray/nditer_pywrap.c index db6c4eb07..8acc7f87f 100644 --- a/numpy/core/src/multiarray/nditer_pywrap.c +++ b/numpy/core/src/multiarray/nditer_pywrap.c @@ -1601,8 +1601,8 @@ npyiter_multi_index_set( for (idim = 0; idim < ndim; ++idim) { PyObject *v = PySequence_GetItem(value, idim); multi_index[idim] = PyLong_AsLong(v); + Py_DECREF(v); if (error_converting(multi_index[idim])) { - Py_XDECREF(v); return -1; } } diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index 1b47e74ac..60a315f6e 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -5228,6 +5228,7 @@ PyUFunc_RegisterLoopForDescr(PyUFuncObject *ufunc, arg_typenums = PyArray_malloc(ufunc->nargs * sizeof(int)); if (arg_typenums == NULL) { + Py_DECREF(key); PyErr_NoMemory(); return -1; } @@ -5365,6 +5366,7 @@ PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc, /* Get entry for this user-defined type*/ cobj = PyDict_GetItemWithError(ufunc->userloops, key); if (cobj == NULL && PyErr_Occurred()) { + Py_DECREF(key); return 0; } /* If it's not there, then make one and return. */ diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py index b44343c57..adcf921f6 100644 --- a/numpy/core/tests/test_nditer.py +++ b/numpy/core/tests/test_nditer.py @@ -185,6 +185,29 @@ def test_iter_c_or_f_order(): assert_equal([x for x in i], aview.swapaxes(0, 1).ravel(order='A')) +def test_nditer_multi_index_set(): + # Test the multi_index set + a = np.arange(6).reshape(2, 3) + it = np.nditer(a, flags=['multi_index']) + + # Removes the iteration on two first elements of a[0] + it.multi_index = (0, 2,) + + assert_equal([i for i in it], [2, 3, 4, 5]) + +@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") +def test_nditer_multi_index_set_refcount(): + # Test if the reference count on index variable is decreased + + index = 0 + i = np.nditer(np.array([111, 222, 333, 444]), flags=['multi_index']) + + start_count = sys.getrefcount(index) + i.multi_index = (index,) + end_count = sys.getrefcount(index) + + assert_equal(start_count, end_count) + def test_iter_best_order_multi_index_1d(): # The multi-indices should be correct with any reordering |