diff options
| -rw-r--r-- | numpy/core/src/multiarray/iterators.c | 50 | ||||
| -rw-r--r-- | numpy/core/tests/test_multiarray.py | 16 |
2 files changed, 41 insertions, 25 deletions
diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c index a3bc8e742..62a057538 100644 --- a/numpy/core/src/multiarray/iterators.c +++ b/numpy/core/src/multiarray/iterators.c @@ -539,6 +539,7 @@ iter_subscript(PyArrayIterObject *self, PyObject *ind) char *dptr; int size; PyObject *obj = NULL; + PyObject *new; PyArray_CopySwapFunc *copyswap; if (ind == Py_Ellipsis) { @@ -640,35 +641,34 @@ iter_subscript(PyArrayIterObject *self, PyObject *ind) obj = ind; } - if (PyArray_Check(obj)) { - /* Check for Boolean object */ - if (PyArray_TYPE((PyArrayObject *)obj) == NPY_BOOL) { - ret = iter_subscript_Bool(self, (PyArrayObject *)obj); - Py_DECREF(indtype); - } - /* Check for integer array */ - else if (PyArray_ISINTEGER((PyArrayObject *)obj)) { - PyObject *new; - new = PyArray_FromAny(obj, indtype, 0, 0, - NPY_ARRAY_FORCECAST | NPY_ARRAY_ALIGNED, NULL); - if (new == NULL) { - goto fail; - } - Py_DECREF(obj); - obj = new; - new = iter_subscript_int(self, (PyArrayObject *)obj); - Py_DECREF(obj); - return new; - } - else { - goto fail; - } + /* Any remaining valid input is an array or has been turned into one */ + if (!PyArray_Check(obj)) { + goto fail; + } + + /* Check for Boolean array */ + if (PyArray_TYPE((PyArrayObject *)obj) == NPY_BOOL) { + ret = iter_subscript_Bool(self, (PyArrayObject *)obj); + Py_DECREF(indtype); Py_DECREF(obj); return (PyObject *)ret; } - else { - Py_DECREF(indtype); + + /* Only integer arrays left */ + if (!PyArray_ISINTEGER((PyArrayObject *)obj)) { + goto fail; + } + + Py_INCREF(indtype); + new = PyArray_FromAny(obj, indtype, 0, 0, + NPY_ARRAY_FORCECAST | NPY_ARRAY_ALIGNED, NULL); + if (new == NULL) { + goto fail; } + Py_DECREF(indtype); + Py_DECREF(obj); + Py_SETREF(new, iter_subscript_int(self, (PyArrayObject *)new)); + return new; fail: diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py index ae2fd3cf4..c45029599 100644 --- a/numpy/core/tests/test_multiarray.py +++ b/numpy/core/tests/test_multiarray.py @@ -4893,6 +4893,22 @@ class TestFlat(object): assert_(e.flags.writebackifcopy is False) assert_(f.flags.writebackifcopy is False) + @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") + def test_refcount(self): + # includes regression test for reference count error gh-13165 + inds = [np.intp(0), np.array([True]*self.a.size), np.array([0]), None] + indtype = np.dtype(np.intp) + rc_indtype = sys.getrefcount(indtype) + for ind in inds: + rc_ind = sys.getrefcount(ind) + for _ in range(100): + try: + self.a.flat[ind] + except IndexError: + pass + assert_(abs(sys.getrefcount(ind) - rc_ind) < 50) + assert_(abs(sys.getrefcount(indtype) - rc_indtype) < 50) + class TestResize(object): def test_basic(self): |