BUG: Incref items in np.take on error as they are decrefed later

When take fails during copying due to out of bound indices, then the already
copied items will be decref'd on array destruction. To avoid that, as well
as possible overlapping or already initialized object arrays, decref and
incref during the copy operation itself. Note that all basic types but
object use their own fasttake, so this does not change anything for them.

numpy/core/src/multiarray/item_selection.c

@@ -31,10 +31,11 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
     PyArray_Descr *dtype;
     PyArray_FastTakeFunc *func;
     PyArrayObject *obj = NULL, *self, *indices;
-    npy_intp nd, i, j, n, m, max_item, tmp, chunk, nelem;
+    npy_intp nd, i, j, n, m, k, max_item, tmp, chunk, itemsize, nelem;
     npy_intp shape[NPY_MAXDIMS];
-    char *src, *dest;
+    char *src, *dest, *tmp_src;
     int err;
+    npy_bool needs_refcounting;
 
     indices = NULL;
     self = (PyArrayObject *)PyArray_CheckAxis(self0, &axis,
@@ -110,9 +111,11 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
 
     max_item = PyArray_DIMS(self)[axis];
     nelem = chunk;
-    chunk = chunk * PyArray_DESCR(obj)->elsize;
+    itemsize = PyArray_ITEMSIZE(obj);
+    chunk = chunk * itemsize;
     src = PyArray_DATA(self);
     dest = PyArray_DATA(obj);
+    needs_refcounting = PyDataType_REFCHK(PyArray_DESCR(self));
 
     func = PyArray_DESCR(self)->f->fasttake;
     if (func == NULL) {
@@ -124,8 +127,20 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
                     if (check_and_adjust_index(&tmp, max_item, axis) < 0) {
                         goto fail;
                     }
-                    memmove(dest, src + tmp*chunk, chunk);
-                    dest += chunk;
+                    tmp_src = src + tmp * chunk;
+                    if (needs_refcounting) {
+                        for (k=0; k < nelem; k++) {
+                            PyArray_Item_INCREF(tmp_src, PyArray_DESCR(self));
+                            PyArray_Item_XDECREF(dest, PyArray_DESCR(self));
+                            memmove(dest, tmp_src, itemsize);
+                            dest += itemsize;
+                            tmp_src += itemsize;
+                        }
+                    }
+                    else {
+                        memmove(dest, tmp_src, chunk);
+                        dest += chunk;
+                    }
                 }
                 src += chunk*max_item;
             }
@@ -144,8 +159,20 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
                             tmp -= max_item;
                         }
                     }
-                    memmove(dest, src + tmp*chunk, chunk);
-                    dest += chunk;
+                    tmp_src = src + tmp * chunk;
+                    if (needs_refcounting) {
+                        for (k=0; k < nelem; k++) {
+                            PyArray_Item_INCREF(tmp_src, PyArray_DESCR(self));
+                            PyArray_Item_XDECREF(dest, PyArray_DESCR(self));
+                            memmove(dest, tmp_src, itemsize);
+                            dest += itemsize;
+                            tmp_src += itemsize;
+                        }
+                    }
+                    else {
+                        memmove(dest, tmp_src, chunk);
+                        dest += chunk;
+                    }
                 }
                 src += chunk*max_item;
             }
@@ -160,8 +187,20 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
                     else if (tmp >= max_item) {
                         tmp = max_item - 1;
                     }
-                    memmove(dest, src+tmp*chunk, chunk);
-                    dest += chunk;
+                    tmp_src = src + tmp * chunk;
+                    if (needs_refcounting) {
+                        for (k=0; k < nelem; k++) {
+                            PyArray_Item_INCREF(tmp_src, PyArray_DESCR(self));
+                            PyArray_Item_XDECREF(dest, PyArray_DESCR(self));
+                            memmove(dest, tmp_src, itemsize);
+                            dest += itemsize;
+                            tmp_src += itemsize;
+                        }
+                    }
+                    else {
+                        memmove(dest, tmp_src, chunk);
+                        dest += chunk;
+                    }
                 }
                 src += chunk*max_item;
             }
@@ -176,7 +215,6 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
         }
     }
 
-    PyArray_INCREF(obj);
     Py_XDECREF(indices);
     Py_XDECREF(self);
     if (out != NULL && out != obj) {

numpy/core/tests/test_item_selection.py

@@ -4,7 +4,7 @@
 
 def test_take():
     a = [[1, 2], [3, 4]]
-    a_str = [['1','2'],['3','4']]
+    a_str = [[b'1', b'2'],[b'3', b'4']]
     modes = ['raise', 'wrap', 'clip']
     indices = [-1, 4]
     index_arrays = [np.empty(0, dtype=np.intp),
@@ -15,10 +15,12 @@ def test_take():
     real_indices['wrap'] = {-1:1, 4:0}
     real_indices['clip'] = {-1:0, 4:1}
     # Currently all types but object, use the same function generation.
-    # So it should not be necessary to test all, but the code does support it.
-    types = np.int, np.object
+    # So it should not be necessary to test all. However test also a non
+    # refcounted struct on top of object.
+    types = np.int, np.object, np.dtype([('', 'i', 2)])
     for t in types:
-        ta = np.array(a if issubclass(t, np.number) else a_str, dtype=t)
+        # ta works, even if the array may be odd if buffer interface is used
+        ta = np.array(a if np.issubdtype(t, np.number) else a_str, dtype=t)
         tresult = list(ta.T.copy())
         for index_array in index_arrays:
             if index_array.size != 0:

numpy/core/tests/test_regression.py

@@ -626,6 +626,17 @@ def test_take_output(self, level=rlevel):
         np.take(x,[0,2],axis=1,out=b)
         assert_array_equal(a,b)
 
+    def test_take_object_fail(self):
+        # Issue gh-3001
+        d = 123.
+        a = np.array([d, 1], dtype=object)
+        ref_d = sys.getrefcount(d)
+        try:
+            a.take([0, 100])
+        except IndexError:
+            pass
+        assert_(ref_d == sys.getrefcount(d))
+
     def test_array_str_64bit(self, level=rlevel):
         """Ticket #501"""
         s = np.array([1, np.nan],dtype=np.float64)