Issue #26015: Added new tests for pickling iterators of mutable sequences.

5 files changed, 214 insertions, 50 deletions

diff --git a/Lib/test/test_array.py b/Lib/test/test_array.py
index 10d9946..482526e 100644
--- a/Lib/test/test_array.py
+++ b/Lib/test/test_array.py
@@ -284,19 +284,42 @@ class BaseTest:
             self.assertEqual(type(a), type(b))
 
     def test_iterator_pickle(self):
-        data = array.array(self.typecode, self.example)
+        orig = array.array(self.typecode, self.example)
+        data = list(orig)
+        data2 = data[::-1]
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
-            orgit = iter(data)
-            d = pickle.dumps(orgit, proto)
-            it = pickle.loads(d)
-            self.assertEqual(type(orgit), type(it))
-            self.assertEqual(list(it), list(data))
-
-            if len(data):
-                it = pickle.loads(d)
-                next(it)
-                d = pickle.dumps(it, proto)
-                self.assertEqual(list(it), list(data)[1:])
+            # initial iterator
+            itorig = iter(orig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a.fromlist(data2)
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data + data2)
+
+            # running iterator
+            next(itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a.fromlist(data2)
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data[1:] + data2)
+
+            # empty iterator
+            for i in range(1, len(data)):
+                next(itorig)
+            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)
+
+            # exhausted iterator
+            self.assertRaises(StopIteration, next, itorig)
+            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)
 
     def test_insert(self):
         a = array.array(self.typecode, self.example)
diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
index e11da5a..5d0351b 100644
--- a/Lib/test/test_bytes.py
+++ b/Lib/test/test_bytes.py
@@ -609,10 +609,9 @@ class BaseBytesTest:
                 self.assertEqual(list(it), data)
 
                 it = pickle.loads(d)
-                try:
-                    next(it)
-                except StopIteration:
+                if not b:
                     continue
+                next(it)
                 d = pickle.dumps(it, proto)
                 it = pickle.loads(d)
                 self.assertEqual(list(it), data[1:])
@@ -1379,6 +1378,43 @@ class ByteArrayTest(BaseBytesTest, unittest.TestCase):
         from _testcapi import getbuffer_with_null_view
         self.assertRaises(BufferError, getbuffer_with_null_view, bytearray())
 
+    def test_iterator_pickling2(self):
+        orig = bytearray(b'abc')
+        data = list(b'qwerty')
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+            # initial iterator
+            itorig = iter(orig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, b = pickle.loads(d)
+            b[:] = data
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data)
+
+            # running iterator
+            next(itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, b = pickle.loads(d)
+            b[:] = data
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data[1:])
+
+            # empty iterator
+            for i in range(1, len(orig)):
+                next(itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, b = pickle.loads(d)
+            b[:] = data
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data[len(orig):])
+
+            # exhausted iterator
+            self.assertRaises(StopIteration, next, itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, b = pickle.loads(d)
+            b[:] = data
+            self.assertEqual(list(it), [])
+
+
 class AssortedBytesTest(unittest.TestCase):
     #
     # Test various combinations of bytes and bytearray
diff --git a/Lib/test/test_deque.py b/Lib/test/test_deque.py
index 634d71f..29b44f6 100644
--- a/Lib/test/test_deque.py
+++ b/Lib/test/test_deque.py
@@ -640,18 +640,45 @@ class TestBasic(unittest.TestCase):
                 self.assertEqual(e.maxlen, d.maxlen)
 
     def test_iterator_pickle(self):
-        data = deque(range(200))
+        orig = deque(range(200))
+        data = [i*1.01 for i in orig]
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
-            it = itorg = iter(data)
-            d = pickle.dumps(it, proto)
-            it = pickle.loads(d)
-            self.assertEqual(type(itorg), type(it))
-            self.assertEqual(list(it), list(data))
-
-            it = pickle.loads(d)
-            next(it)
-            d = pickle.dumps(it, proto)
-            self.assertEqual(list(it), list(data)[1:])
+            # initial iterator
+            itorg = iter(orig)
+            dump = pickle.dumps((itorg, orig), proto)
+            it, d = pickle.loads(dump)
+            for i, x in enumerate(data):
+                d[i] = x
+            self.assertEqual(type(it), type(itorg))
+            self.assertEqual(list(it), data)
+
+            # running iterator
+            next(itorg)
+            dump = pickle.dumps((itorg, orig), proto)
+            it, d = pickle.loads(dump)
+            for i, x in enumerate(data):
+                d[i] = x
+            self.assertEqual(type(it), type(itorg))
+            self.assertEqual(list(it), data[1:])
+
+            # empty iterator
+            for i in range(1, len(data)):
+                next(itorg)
+            dump = pickle.dumps((itorg, orig), proto)
+            it, d = pickle.loads(dump)
+            for i, x in enumerate(data):
+                d[i] = x
+            self.assertEqual(type(it), type(itorg))
+            self.assertEqual(list(it), [])
+
+            # exhausted iterator
+            self.assertRaises(StopIteration, next, itorg)
+            dump = pickle.dumps((itorg, orig), proto)
+            it, d = pickle.loads(dump)
+            for i, x in enumerate(data):
+                d[i] = x
+            self.assertEqual(type(it), type(itorg))
+            self.assertEqual(list(it), [])
 
     def test_deepcopy(self):
         mut = [10]
diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
index abc408f..56e21f8 100644
--- a/Lib/test/test_iter.py
+++ b/Lib/test/test_iter.py
@@ -153,6 +153,42 @@ class TestCase(unittest.TestCase):
     def test_seq_class_iter(self):
         self.check_iterator(iter(SequenceClass(10)), list(range(10)))
 
+    def test_mutating_seq_class_iter_pickle(self):
+        orig = SequenceClass(5)
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+            # initial iterator
+            itorig = iter(orig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, seq = pickle.loads(d)
+            seq.n = 7
+            self.assertIs(type(it), type(itorig))
+            self.assertEqual(list(it), list(range(7)))
+
+            # running iterator
+            next(itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, seq = pickle.loads(d)
+            seq.n = 7
+            self.assertIs(type(it), type(itorig))
+            self.assertEqual(list(it), list(range(1, 7)))
+
+            # empty iterator
+            for i in range(1, 5):
+                next(itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, seq = pickle.loads(d)
+            seq.n = 7
+            self.assertIs(type(it), type(itorig))
+            self.assertEqual(list(it), list(range(5, 7)))
+
+            # exhausted iterator
+            self.assertRaises(StopIteration, next, itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, seq = pickle.loads(d)
+            seq.n = 7
+            self.assertTrue(isinstance(it, collections.abc.Iterator))
+            self.assertEqual(list(it), [])
+
     # Test a new_style class with __iter__ but no next() method
     def test_new_style_iter_class(self):
         class IterClass(object):
diff --git a/Lib/test/test_list.py b/Lib/test/test_list.py
index 750d6cf..8f82ab5 100644
--- a/Lib/test/test_list.py
+++ b/Lib/test/test_list.py
@@ -72,34 +72,76 @@ class ListTest(list_tests.CommonTest):
         check(1000000)
 
     def test_iterator_pickle(self):
-        # Userlist iterators don't support pickling yet since
-        # they are based on generators.
-        data = self.type2test([4, 5, 6, 7])
+        orig = self.type2test([4, 5, 6, 7])
+        data = [10, 11, 12, 13, 14, 15]
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
-            it = itorg = iter(data)
-            d = pickle.dumps(it, proto)
-            it = pickle.loads(d)
-            self.assertEqual(type(itorg), type(it))
-            self.assertEqual(self.type2test(it), self.type2test(data))
-
-            it = pickle.loads(d)
-            next(it)
-            d = pickle.dumps(it, proto)
-            self.assertEqual(self.type2test(it), self.type2test(data)[1:])
+            # initial iterator
+            itorig = iter(orig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a[:] = data
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data)
+
+            # running iterator
+            next(itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a[:] = data
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data[1:])
+
+            # empty iterator
+            for i in range(1, len(orig)):
+                next(itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a[:] = data
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data[len(orig):])
+
+            # exhausted iterator
+            self.assertRaises(StopIteration, next, itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a[:] = data
+            self.assertEqual(list(it), [])
 
     def test_reversed_pickle(self):
-        data = self.type2test([4, 5, 6, 7])
+        orig = self.type2test([4, 5, 6, 7])
+        data = [10, 11, 12, 13, 14, 15]
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
-            it = itorg = reversed(data)
-            d = pickle.dumps(it, proto)
-            it = pickle.loads(d)
-            self.assertEqual(type(itorg), type(it))
-            self.assertEqual(self.type2test(it), self.type2test(reversed(data)))
-
-            it = pickle.loads(d)
-            next(it)
-            d = pickle.dumps(it, proto)
-            self.assertEqual(self.type2test(it), self.type2test(reversed(data))[1:])
+            # initial iterator
+            itorig = reversed(orig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a[:] = data
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data[len(orig)-1::-1])
+
+            # running iterator
+            next(itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a[:] = data
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), data[len(orig)-2::-1])
+
+            # empty iterator
+            for i in range(1, len(orig)):
+                next(itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a[:] = data
+            self.assertEqual(type(it), type(itorig))
+            self.assertEqual(list(it), [])
+
+            # exhausted iterator
+            self.assertRaises(StopIteration, next, itorig)
+            d = pickle.dumps((itorig, orig), proto)
+            it, a = pickle.loads(d)
+            a[:] = data
+            self.assertEqual(list(it), [])
 
     def test_no_comdat_folding(self):
         # Issue 8847: In the PGO build, the MSVC linker's COMDAT folding