diff --git a/numpy/core/src/multiarray/nditer_pywrap.c b/numpy/core/src/multiarray/nditer_pywrap.c index 7698ae43d07e..73df962e4520 100644 --- a/numpy/core/src/multiarray/nditer_pywrap.c +++ b/numpy/core/src/multiarray/nditer_pywrap.c @@ -1595,8 +1595,8 @@ npyiter_multi_index_set(NewNpyArrayIterObject *self, PyObject *value) 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/tests/test_nditer.py b/numpy/core/tests/test_nditer.py index b44343c5755c..adcf921f6016 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