ENH: iter: Use the new code to do buffer decrefs

1 files changed, 21 insertions, 12 deletions

diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index 85eaee7fb..99c6decd9 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -4250,6 +4250,19 @@ npyiter_allocate_transfer_functions(NpyIter *iter)
                 writetransferfn[iiter] = stransfer;
                 writetransferdata[iiter] = transferdata;
             }
+            /* If no write back but there are references make a decref fn */
+            else if (PyDataType_REFCHK(op_dtype[iiter])) {
+                if (PyArray_GetDecSrcRefTransferFunction(
+                                        (flags&NPY_OP_ITFLAG_ALIGNED) != 0,
+                                        op_dtype[iiter]->elsize,
+                                        op_dtype[iiter],
+                                        &stransfer,
+                                        &transferdata) != NPY_SUCCEED) {
+                    goto fail;
+                }
+                writetransferfn[iiter] = stransfer;
+                writetransferdata[iiter] = transferdata;
+            }
             else {
                 writetransferfn[iiter] = NULL;
             }
@@ -4472,21 +4485,17 @@ npyiter_copy_from_buffers(NpyIter *iter)
                         transferdata);
             }
         }
-        /* If there's no copy back, we may have to decrement refs */
-        else if (PyDataType_REFCHK(dtypes[iiter])) {
+        /* If there's no copy back, we may have to decrement refs.  In
+         * this case, the transfer function has a 'decsrcref' transfer
+         * function, so we can use it to do the decrement.
+         */
+        else if (stransfer != NULL) {
             /* Decrement refs only if the pointer was pointing to the buffer */
             npy_intp delta = (ptrs[iiter] - buffer);
             if (0 <= delta && delta <= transfersize*dtypes[iiter]->elsize) {
-                npy_intp i, size = NBF_SIZE(bufferdata),
-                        itemsize = dtypes[iiter]->elsize;
-
-                for (i = 0; i < size; ++i) {
-                    PyArray_Item_XDECREF(buffer, dtypes[iiter]);
-                    buffer += itemsize;
-                }
-
-                /* Also set the values to zero for safety */
-                memset(buffers[iiter], 0, itemsize*size);
+                stransfer(NULL, 0, buffer, dtypes[iiter]->elsize,
+                            transfersize, dtypes[iiter]->elsize,
+                            transferdata);
             }
         }
     }