BUG: it.close() disallows access to iterator, fixes #10950

2 files changed, 28 insertions, 17 deletions

diff --git a/numpy/core/src/multiarray/nditer_pywrap.c b/numpy/core/src/multiarray/nditer_pywrap.c
index d36be61f5..d00ecc984 100644
--- a/numpy/core/src/multiarray/nditer_pywrap.c
+++ b/numpy/core/src/multiarray/nditer_pywrap.c
@@ -20,7 +20,7 @@
 
 typedef struct NewNpyArrayIterObject_tag NewNpyArrayIterObject;
 
-enum NPYITER_CONTEXT {CONTEXT_NOTENTERED, CONTEXT_INSIDE, CONTEXT_EXITED};
+enum NPYITER_CONTEXT {CONTEXT_NOTENTERED, CONTEXT_INSIDE, CLOSED};
 
 struct NewNpyArrayIterObject_tag {
     PyObject_HEAD
@@ -1419,7 +1419,7 @@ static PyObject *npyiter_value_get(NewNpyArrayIterObject *self)
         ret = npyiter_seq_item(self, 0);
     }
     else {
-        if (self->managed == CONTEXT_EXITED) {
+        if (self->managed == CLOSED) {
             PyErr_SetString(PyExc_ValueError,
                     "Iterator is closed");
             return NULL;
@@ -1454,7 +1454,7 @@ static PyObject *npyiter_operands_get(NewNpyArrayIterObject *self)
                 "Iterator is invalid");
         return NULL;
     }
-    if (self->managed == CONTEXT_EXITED) {
+    if (self->managed == CLOSED) {
         PyErr_SetString(PyExc_ValueError,
                 "Iterator is closed");
         return NULL;
@@ -1489,7 +1489,7 @@ static PyObject *npyiter_itviews_get(NewNpyArrayIterObject *self)
         return NULL;
     }
 
-    if (self->managed == CONTEXT_EXITED) {
+    if (self->managed == CLOSED) {
         PyErr_SetString(PyExc_ValueError,
                 "Iterator is closed");
         return NULL;
@@ -1517,7 +1517,8 @@ static PyObject *npyiter_itviews_get(NewNpyArrayIterObject *self)
 static PyObject *
 npyiter_next(NewNpyArrayIterObject *self)
 {
-    if (self->iter == NULL || self->iternext == NULL || self->finished) {
+    if (self->iter == NULL || self->iternext == NULL ||
+                self->finished || (self->managed == CLOSED)) {
         return NULL;
     }
 
@@ -1912,7 +1913,7 @@ static PyObject *npyiter_dtypes_get(NewNpyArrayIterObject *self)
         return NULL;
     }
 
-    if (self->managed == CONTEXT_EXITED) {
+    if (self->managed == CLOSED) {
         PyErr_SetString(PyExc_ValueError,
                 "Iterator is closed");
         return NULL;
@@ -2014,7 +2015,7 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
         return NULL;
     }
 
-    if (self->managed == CONTEXT_EXITED) {
+    if (self->managed == CLOSED) {
         PyErr_SetString(PyExc_ValueError,
                 "Iterator is closed");
         return NULL;
@@ -2104,7 +2105,7 @@ npyiter_seq_slice(NewNpyArrayIterObject *self,
         return NULL;
     }
 
-    if (self->managed == CONTEXT_EXITED) {
+    if (self->managed == CLOSED) {
         PyErr_SetString(PyExc_ValueError,
                 "Iterator is closed");
         return NULL;
@@ -2170,7 +2171,7 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
         return -1;
     }
 
-    if (self->managed == CONTEXT_EXITED) {
+    if (self->managed == CLOSED) {
         PyErr_SetString(PyExc_ValueError,
                 "Iterator is closed");
         return -1;
@@ -2250,7 +2251,7 @@ npyiter_seq_ass_slice(NewNpyArrayIterObject *self, Py_ssize_t ilow,
         return -1;
     }
 
-    if (self->managed == CONTEXT_EXITED) {
+    if (self->managed == CLOSED) {
         PyErr_SetString(PyExc_ValueError,
                 "Iterator is closed");
         return -1;
@@ -2307,7 +2308,7 @@ npyiter_subscript(NewNpyArrayIterObject *self, PyObject *op)
         return NULL;
     }
 
-    if (self->managed == CONTEXT_EXITED) {
+    if (self->managed == CLOSED) {
         PyErr_SetString(PyExc_ValueError,
                 "Iterator is closed");
         return NULL;
@@ -2362,7 +2363,7 @@ npyiter_ass_subscript(NewNpyArrayIterObject *self, PyObject *op,
         return -1;
     }
 
-    if (self->managed == CONTEXT_EXITED) {
+    if (self->managed == CLOSED) {
         PyErr_SetString(PyExc_ValueError,
                 "Iterator is closed");
         return -1;
@@ -2402,8 +2403,8 @@ npyiter_enter(NewNpyArrayIterObject *self)
         PyErr_SetString(PyExc_RuntimeError, "operation on non-initialized iterator");
         return NULL;
     }
-    if (self->managed == CONTEXT_EXITED) {
-        PyErr_SetString(PyExc_ValueError, "cannot reuse iterator after exit");
+    if (self->managed == CLOSED) {
+        PyErr_SetString(PyExc_ValueError, "cannot reuse closed iterator");
         return NULL;
     }
     self->managed = CONTEXT_INSIDE;
@@ -2420,6 +2421,7 @@ npyiter_close(NewNpyArrayIterObject *self)
         Py_RETURN_NONE;
     }
     ret = NpyIter_Close(iter);
+    self->managed = CLOSED;
     if (ret < 0) {
         return NULL;
     }
@@ -2429,7 +2431,7 @@ npyiter_close(NewNpyArrayIterObject *self)
 static PyObject *
 npyiter_exit(NewNpyArrayIterObject *self, PyObject *args)
 {
-    self->managed = CONTEXT_EXITED;
+    self->managed = CLOSED;
     /* even if called via exception handling, writeback any data */
     return npyiter_close(self);
 }
diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py
index bc9456536..77c26eacf 100644
--- a/numpy/core/tests/test_nditer.py
+++ b/numpy/core/tests/test_nditer.py
@@ -2847,7 +2847,7 @@ def test_writebacks():
     enter = it.__enter__
     assert_raises(ValueError, enter)
 
-def test_close():
+def test_close_equivalent():
     ''' using a context amanger and using nditer.close are equivalent
     '''
     def add_close(x, y, out=None):
@@ -2856,8 +2856,10 @@ def test_close():
                     [['readonly'], ['readonly'], ['writeonly','allocate']])
         for (a, b, c) in it:
             addop(a, b, out=c)
+        ret = it.operands[2]
         it.close()
-        return it.operands[2]
+        return ret
+
     def add_context(x, y, out=None):
         addop = np.add
         it = np.nditer([x, y, out], [],
@@ -2871,6 +2873,13 @@ def test_close():
     z = add_context(range(5), range(5))
     assert_equal(z, range(0, 10, 2))
 
+def test_close_raises():
+    it = np.nditer(np.arange(3))
+    assert_equal (next(it), 0)
+    it.close()
+    assert_raises(StopIteration, next, it)
+    assert_raises(ValueError, getattr, it, 'operands')
+
 def test_warn_noclose():
     a = np.arange(6, dtype='f4')
     au = a.byteswap().newbyteorder()