diff options
Diffstat (limited to 'numpy')
-rw-r--r-- | numpy/core/src/multiarray/new_iterator.c.src | 33 |
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); } } } |