BUG: core: fix refcount error in PyArray_Take (#939)

2 files changed, 15 insertions, 0 deletions

diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
index 398acfe71..224e11e3b 100644
--- a/numpy/core/src/multiarray/item_selection.c
+++ b/numpy/core/src/multiarray/item_selection.c
@@ -43,6 +43,7 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
                                                          PyArray_INTP,
                                                          1, 0);
     if (indices == NULL) {
+        Py_XINCREF(ret);
         goto fail;
     }
     n = m = chunk = 1;
diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py
index 1257a54c2..5959d6dd9 100644
--- a/numpy/core/tests/test_regression.py
+++ b/numpy/core/tests/test_regression.py
@@ -1359,5 +1359,19 @@ class TestRegression(TestCase):
         y = np.add(x, x, x)
         assert_equal(id(x), id(y))
 
+    def test_take_refcount(self):
+        # ticket #939
+        a = np.arange(16, dtype=np.float)
+        a.shape = (4,4)
+        lut = np.ones((5 + 3, 4), np.float)
+        rgba = np.empty(shape=a.shape + (4,), dtype=lut.dtype)
+        c1 = sys.getrefcount(rgba)
+        try:
+            lut.take(a, axis=0, mode='clip', out=rgba)
+        except TypeError:
+            pass
+        c2 = sys.getrefcount(rgba)
+        assert_equal(c1, c2)
+
 if __name__ == "__main__":
     run_module_suite()