TST: Add a test for nditer write masked with references (gh-18374)

* TST: Add a test for nditer write masked with references

This code is almost impossible to hit, but in theory, a masked
iterator can contain references which are buffered.
The buffer than would use the `move_reference` mechanism,
Even though not all elements are copied, all references in the
buffer must be cleared in the process.

* BUG: fix test for non-refcount

Co-authored-by: Matti Picus <matti.picus@gmail.com>

2 files changed, 32 insertions, 0 deletions

diff --git a/numpy/core/src/multiarray/dtype_transfer.c b/numpy/core/src/multiarray/dtype_transfer.c
index 9b8e5f32f..fa17caced 100644
--- a/numpy/core/src/multiarray/dtype_transfer.c
+++ b/numpy/core/src/multiarray/dtype_transfer.c
@@ -4625,6 +4625,7 @@ PyArray_GetMaskedDTypeTransferFunction(int aligned,
     /* TODO: Special case some important cases so they're fast */
 
     /* Fall back to wrapping a non-masked transfer function */
+    assert(dst_dtype != NULL);
     if (PyArray_GetDTypeTransferFunction(aligned,
                                 src_stride, dst_stride,
                                 src_dtype, dst_dtype,
diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py
index 4271d2d96..411095199 100644
--- a/numpy/core/tests/test_nditer.py
+++ b/numpy/core/tests/test_nditer.py
@@ -2746,6 +2746,37 @@ def test_iter_writemasked():
     # were copied back
     assert_equal(a, [3, 3, 2.5])
 
+def test_iter_writemasked_decref():
+    # force casting (to make it interesting) by using a structured dtype.
+    arr = np.arange(10000).astype(">i,O")
+    original = arr.copy()
+    mask = np.random.randint(0, 2, size=10000).astype(bool)
+
+    it = np.nditer([arr, mask], ['buffered', "refs_ok"],
+                   [['readwrite', 'writemasked'],
+                    ['readonly', 'arraymask']],
+                   op_dtypes=["<i,O", "?"])
+    singleton = object()
+    if HAS_REFCOUNT:
+        count = sys.getrefcount(singleton)
+    for buf, mask_buf in it:
+        buf[...] = (3, singleton)
+
+    del buf, mask_buf, it   # delete everything to ensure corrrect cleanup
+
+    if HAS_REFCOUNT:
+        # The buffer would have included additional items, they must be
+        # cleared correctly:
+        assert sys.getrefcount(singleton) - count == np.count_nonzero(mask)
+
+    assert_array_equal(arr[~mask], original[~mask])
+    assert (arr[mask] == np.array((3, singleton), arr.dtype)).all()
+    del arr
+
+    if HAS_REFCOUNT:
+        assert sys.getrefcount(singleton) == count
+
+
 def test_iter_non_writable_attribute_deletion():
     it = np.nditer(np.ones(2))
     attr = ["value", "shape", "operands", "itviews", "has_delayed_bufalloc",