BUG: Fix decref before incref for in-place accumulate

seberg · jaimefrio · commit e2684d0783c3 · 2016-03-15T08:22:46.000+01:00
closes numpygh-7402
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
@@ -3270,17 +3270,22 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
         NPY_BEGIN_THREADS_NDITER(iter);
 
         do {
-
             dataptr_copy[0] = dataptr[0];
             dataptr_copy[1] = dataptr[1];
             dataptr_copy[2] = dataptr[0];
 
             /* Copy the first element to start the reduction */
             if (otype == NPY_OBJECT) {
+                /*
+                 * Input (dataptr[0]) and output (dataptr[1]) may point
+                 * to the same memory (i.e. np.add.accumulate(a, out=a)).
+                 * In that case need to incref before decref to avoid the
+                 * possibility of the reference count being zero temporarily.
+                 */
+                Py_XINCREF(*(PyObject **)dataptr_copy[1]);
                 Py_XDECREF(*(PyObject **)dataptr_copy[0]);
                 *(PyObject **)dataptr_copy[0] =
                                     *(PyObject **)dataptr_copy[1];
-                Py_XINCREF(*(PyObject **)dataptr_copy[0]);
             }
             else {
                 memcpy(dataptr_copy[0], dataptr_copy[1], itemsize);
@@ -3333,10 +3338,16 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
 
         /* Copy the first element to start the reduction */
         if (otype == NPY_OBJECT) {
+            /*
+             * Input (dataptr[0]) and output (dataptr[1]) may point
+             * to the same memory (i.e. np.add.accumulate(a, out=a, axis=0)).
+             * In that case need to incref before decref to avoid the
+             * possibility of the reference count being zero temporarily.
+             */
+            Py_XINCREF(*(PyObject **)dataptr_copy[1]);
             Py_XDECREF(*(PyObject **)dataptr_copy[0]);
             *(PyObject **)dataptr_copy[0] =
                                 *(PyObject **)dataptr_copy[1];
-            Py_XINCREF(*(PyObject **)dataptr_copy[0]);
         }
         else {
             memcpy(dataptr_copy[0], dataptr_copy[1], itemsize);
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
@@ -649,6 +649,22 @@ def test_object_array_reduction(self):
         assert_equal(np.array([[1]], dtype=object).sum(), 1)
         assert_equal(np.array([[[1, 2]]], dtype=object).sum((0, 1)), [1, 2])
 
+    def test_object_array_accumulate_inplace(self):
+        # Checks that in-place accumulates work, see also gh-7402
+        arr = np.ones(4, dtype=object)
+        arr[:] = [[1] for i in range(4)]
+        # Twice reproduced also for tuples:
+        np.add.accumulate(arr, out=arr)
+        np.add.accumulate(arr, out=arr)
+        assert_array_equal(arr, np.array([[1]*i for i in [1, 3, 6, 10]]))
+
+        # And the same if the axis argument is used
+        arr = np.ones((2, 4), dtype=object)
+        arr[0, :] = [[2] for i in range(4)]
+        np.add.accumulate(arr, out=arr, axis=-1)
+        np.add.accumulate(arr, out=arr, axis=-1)
+        assert_array_equal(arr[0, :], np.array([[2]*i for i in [1, 3, 6, 10]]))
+
     def test_object_scalar_multiply(self):
         # Tickets #2469 and #4482
         arr = np.matrix([1, 2], dtype=object)
