Merge pull request numpy#19396 from defoishugo/fix_numpy.nditer_leak

BUG: fix a numpy.npiter leak in npyiter_multi_index_set

Author: mattip

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;
             }
         }

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