Issue #26492: Exhausted iterator of array.array now conforms with the behavior

serhiy-storchaka · serhiy-storchaka · commit ab0d198c7a6b · 2016-03-30T21:11:16.000+03:00
of iterators of other mutable sequences: it lefts exhausted even if iterated
array is extended.
diff --git a/Lib/test/test_array.py b/Lib/test/test_array.py
@@ -318,8 +318,19 @@ def test_iterator_pickle(self):
             d = pickle.dumps((itorig, orig), proto)
             it, a = pickle.loads(d)
             a.fromlist(data2)
-            self.assertEqual(type(it), type(itorig))
-            self.assertEqual(list(it), data2)
+            self.assertEqual(list(it), [])
+
+    def test_exhausted_iterator(self):
+        a = array.array(self.typecode, self.example)
+        self.assertEqual(list(a), list(self.example))
+        exhit = iter(a)
+        empit = iter(a)
+        for x in exhit:  # exhaust the iterator
+            next(empit)  # not exhausted
+        a.append(self.outside)
+        self.assertEqual(list(exhit), [])
+        self.assertEqual(list(empit), [self.outside])
+        self.assertEqual(list(a), list(self.example) + [self.outside])
 
     def test_insert(self):
         a = array.array(self.typecode, self.example)
@@ -1070,6 +1081,12 @@ def test_obsolete_write_lock(self):
         a = array.array('B', b"")
         self.assertRaises(BufferError, getbuffer_with_null_view, a)
 
+    def test_free_after_iterating(self):
+        support.check_free_after_iterating(self, iter, array.array,
+                                           (self.typecode,))
+        support.check_free_after_iterating(self, reversed, array.array,
+                                           (self.typecode,))
+
 class StringTest(BaseTest):
 
     def test_setitem(self):
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -237,6 +237,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #26492: Exhausted iterator of array.array now conforms with the behavior
+  of iterators of other mutable sequences: it lefts exhausted even if iterated
+  array is extended.
+
 - Issue #26641: doctest.DocFileTest and doctest.testfile() now support
   packages (module splitted into multiple directories) for the package
   parameter.
diff --git a/Modules/arraymodule.c b/Modules/arraymodule.c
@@ -2875,9 +2875,20 @@ array_iter(arrayobject *ao)
 static PyObject *
 arrayiter_next(arrayiterobject *it)
 {
+    arrayobject *ao;
+
+    assert(it != NULL);
     assert(PyArrayIter_Check(it));
-    if (it->index < Py_SIZE(it->ao))
-        return (*it->getitem)(it->ao, it->index++);
+    ao = it->ao;
+    if (ao == NULL) {
+        return NULL;
+    }
+    assert(array_Check(ao));
+    if (it->index < Py_SIZE(ao)) {
+        return (*it->getitem)(ao, it->index++);
+    }
+    it->ao = NULL;
+    Py_DECREF(ao);
     return NULL;
 }
 
@@ -2906,8 +2917,11 @@ static PyObject *
 array_arrayiterator___reduce___impl(arrayiterobject *self)
 /*[clinic end generated code: output=7898a52e8e66e016 input=a062ea1e9951417a]*/
 {
-    return Py_BuildValue("N(O)n", _PyObject_GetBuiltin("iter"),
-                         self->ao, self->index);
+    PyObject *func = _PyObject_GetBuiltin("iter");
+    if (self->ao == NULL) {
+        return Py_BuildValue("N(())", func);
+    }
+    return Py_BuildValue("N(O)n", func, self->ao, self->index);
 }
 
 /*[clinic input]
