Issue python#14288: Serialization support for builtin iterators.

Include/iterobject.h

@@ -18,6 +18,8 @@ PyAPI_FUNC(PyObject *) PySeqIter_New(PyObject *);
 
 PyAPI_FUNC(PyObject *) PyCallIter_New(PyObject *, PyObject *);
 
+PyAPI_FUNC(PyObject *) _PyIter_GetBuiltin(const char *iter);
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/seq_tests.py

@@ -4,6 +4,7 @@
 
 import unittest
 import sys
+import pickle
 
 # Various iterables
 # This is used for checking the constructor (here and in test_deque.py)
@@ -388,3 +389,9 @@ def __eq__(self, other):
         self.assertEqual(a.index(0, -4*sys.maxsize, 4*sys.maxsize), 2)
         self.assertRaises(ValueError, a.index, 0, 4*sys.maxsize,-4*sys.maxsize)
         self.assertRaises(ValueError, a.index, 2, 0, -10)
+
+    def test_pickle(self):
+        lst = self.type2test([4, 5, 6, 7])
+        lst2 = pickle.loads(pickle.dumps(lst))
+        self.assertEqual(lst2, lst)
+        self.assertNotEqual(id(lst2), id(lst))

Lib/test/test_array.py

@@ -285,6 +285,20 @@ def test_pickle_for_empty_array(self):
             self.assertEqual(a.x, b.x)
             self.assertEqual(type(a), type(b))
 
+    def test_iterator_pickle(self):
+        data = array.array(self.typecode, self.example)
+        orgit = iter(data)
+        d = pickle.dumps(orgit)
+        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)
+            self.assertEqual(list(it), list(data)[1:])
+
     def test_insert(self):
         a = array.array(self.typecode, self.example)
         a.insert(0, self.example[0])

Lib/test/test_builtin.py

@@ -14,6 +14,7 @@
 import traceback
 from test.support import TESTFN, unlink,  run_unittest, check_warnings
 from operator import neg
+import pickle
 try:
     import pty, signal
 except ImportError:
@@ -110,7 +111,30 @@ class TestFailingIter:
     def __iter__(self):
         raise RuntimeError
 
+def filter_char(arg):
+    return ord(arg) > ord("d")
+
+def map_char(arg):
+    return chr(ord(arg)+1)
+
 class BuiltinTest(unittest.TestCase):
+    # Helper to check picklability
+    def check_iter_pickle(self, it, seq):
+        itorg = it
+        d = pickle.dumps(it)
+        it = pickle.loads(d)
+        self.assertEqual(type(itorg), type(it))
+        self.assertEqual(list(it), seq)
+
+        #test the iterator after dropping one from it
+        it = pickle.loads(d)
+        try:
+            next(it)
+        except StopIteration:
+            return
+        d = pickle.dumps(it)
+        it = pickle.loads(d)
+        self.assertEqual(list(it), seq[1:])
 
     def test_import(self):
         __import__('sys')
@@ -566,6 +590,11 @@ def badfunc():
         self.assertEqual(list(filter(lambda x: x>=3, (1, 2, 3, 4))), [3, 4])
         self.assertRaises(TypeError, list, filter(42, (1, 2)))
 
+    def test_filter_pickle(self):
+        f1 = filter(filter_char, "abcdeabcde")
+        f2 = filter(filter_char, "abcdeabcde")
+        self.check_iter_pickle(f1, list(f2))
+
     def test_getattr(self):
         self.assertTrue(getattr(sys, 'stdout') is sys.stdout)
         self.assertRaises(TypeError, getattr, sys, 1)
@@ -759,6 +788,11 @@ def badfunc(x):
             raise RuntimeError
         self.assertRaises(RuntimeError, list, map(badfunc, range(5)))
 
+    def test_map_pickle(self):
+        m1 = map(map_char, "Is this the real life?")
+        m2 = map(map_char, "Is this the real life?")
+        self.check_iter_pickle(m1, list(m2))
+
     def test_max(self):
         self.assertEqual(max('123123'), '3')
         self.assertEqual(max(1, 2, 3), 3)
@@ -1300,6 +1334,13 @@ def __getitem__(self, i):
                     return i
         self.assertRaises(ValueError, list, zip(BadSeq(), BadSeq()))
 
+    def test_zip_pickle(self):
+        a = (1, 2, 3)
+        b = (4, 5, 6)
+        t = [(1, 4), (2, 5), (3, 6)]
+        z1 = zip(a, b)
+        self.check_iter_pickle(z1, t)
+
     def test_format(self):
         # Test the basic machinery of the format() builtin.  Don't test
         #  the specifics of the various formatters

Lib/test/test_bytes.py

@@ -518,6 +518,24 @@ def test_pickling(self):
                 q = pickle.loads(ps)
                 self.assertEqual(b, q)
 
+    def test_iterator_pickling(self):
+        for b in b"", b"a", b"abc", b"\xffab\x80", b"\0\0\377\0\0":
+            it = itorg = iter(self.type2test(b))
+            data = list(self.type2test(b))
+            d = pickle.dumps(it)
+            it = pickle.loads(d)
+            self.assertEqual(type(itorg), type(it))
+            self.assertEqual(list(it), data)
+
+            it = pickle.loads(d)
+            try:
+                next(it)
+            except StopIteration:
+                continue
+            d = pickle.dumps(it)
+            it = pickle.loads(d)
+            self.assertEqual(list(it), data[1:])
+
     def test_strip(self):
         b = self.type2test(b'mississippi')
         self.assertEqual(b.strip(b'i'), b'mississipp')

Lib/test/test_deque.py

@@ -471,6 +471,19 @@ def test_pickle(self):
 ##            self.assertNotEqual(id(d), id(e))
 ##            self.assertEqual(id(e), id(e[-1]))
 
+    def test_iterator_pickle(self):
+        data = deque(range(200))
+        it = itorg = iter(data)
+        d = pickle.dumps(it)
+        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)
+        self.assertEqual(list(it), list(data)[1:])
+
     def test_deepcopy(self):
         mut = [10]
         d = deque([mut])

Lib/test/test_dict.py

@@ -2,7 +2,9 @@
 from test import support
 
 import collections, random, string
+import collections.abc
 import gc, weakref
+import pickle
 
 
 class DictTest(unittest.TestCase):
@@ -803,6 +805,58 @@ class MyDict(dict):
             pass
         self._tracked(MyDict())
 
+    def test_iterator_pickling(self):
+        data = {1:"a", 2:"b", 3:"c"}
+        it = iter(data)
+        d = pickle.dumps(it)
+        it = pickle.loads(d)
+        self.assertEqual(sorted(it), sorted(data))
+
+        it = pickle.loads(d)
+        try:
+            drop = next(it)
+        except StopIteration:
+            return
+        d = pickle.dumps(it)
+        it = pickle.loads(d)
+        del data[drop]
+        self.assertEqual(sorted(it), sorted(data))
+
+    def test_itemiterator_pickling(self):
+        data = {1:"a", 2:"b", 3:"c"}
+        # dictviews aren't picklable, only their iterators
+        itorg = iter(data.items())
+        d = pickle.dumps(itorg)
+        it = pickle.loads(d)
+        # note that the type of type of the unpickled iterator
+        # is not necessarily the same as the original.  It is
+        # merely an object supporting the iterator protocol, yielding
+        # the same objects as the original one.
+        # self.assertEqual(type(itorg), type(it))
+        self.assertTrue(isinstance(it, collections.abc.Iterator))
+        self.assertEqual(dict(it), data)
+
+        it = pickle.loads(d)
+        drop = next(it)
+        d = pickle.dumps(it)
+        it = pickle.loads(d)
+        del data[drop[0]]
+        self.assertEqual(dict(it), data)
+
+    def test_valuesiterator_pickling(self):
+        data = {1:"a", 2:"b", 3:"c"}
+        # data.values() isn't picklable, only its iterator
+        it = iter(data.values())
+        d = pickle.dumps(it)
+        it = pickle.loads(d)
+        self.assertEqual(sorted(list(it)), sorted(list(data.values())))
+
+        it = pickle.loads(d)
+        drop = next(it)
+        d = pickle.dumps(it)
+        it = pickle.loads(d)
+        values = list(it) + [drop]
+        self.assertEqual(sorted(values), sorted(list(data.values())))
 
 from test import mapping_tests
 

Lib/test/test_enumerate.py

@@ -1,5 +1,6 @@
 import unittest
 import sys
+import pickle
 
 from test import support
 
@@ -61,7 +62,25 @@ def __init__(self, seqn):
     def __iter__(self):
         return self
 
-class EnumerateTestCase(unittest.TestCase):
+class PickleTest:
+    # Helper to check picklability
+    def check_pickle(self, itorg, seq):
+        d = pickle.dumps(itorg)
+        it = pickle.loads(d)
+        self.assertEqual(type(itorg), type(it))
+        self.assertEqual(list(it), seq)
+
+        it = pickle.loads(d)
+        try:
+            next(it)
+        except StopIteration:
+            self.assertFalse(seq[1:])
+            return
+        d = pickle.dumps(it)
+        it = pickle.loads(d)
+        self.assertEqual(list(it), seq[1:])
+
+class EnumerateTestCase(unittest.TestCase, PickleTest):
 
     enum = enumerate
     seq, res = 'abc', [(0,'a'), (1,'b'), (2,'c')]
@@ -73,6 +92,9 @@ def test_basicfunction(self):
         self.assertEqual(list(self.enum(self.seq)), self.res)
         self.enum.__doc__
 
+    def test_pickle(self):
+        self.check_pickle(self.enum(self.seq), self.res)
+
     def test_getitemseqn(self):
         self.assertEqual(list(self.enum(G(self.seq))), self.res)
         e = self.enum(G(''))
@@ -126,7 +148,7 @@ class TestBig(EnumerateTestCase):
     seq = range(10,20000,2)
     res = list(zip(range(20000), seq))
 
-class TestReversed(unittest.TestCase):
+class TestReversed(unittest.TestCase, PickleTest):
 
     def test_simple(self):
         class A:
@@ -212,6 +234,10 @@ def __len__(self): return 2
         ngi = NoGetItem()
         self.assertRaises(TypeError, reversed, ngi)
 
+    def test_pickle(self):
+        for data in 'abc', range(5), tuple(enumerate('abc')), range(1,17,5):
+            self.check_pickle(reversed(data), list(data)[::-1])
+
 
 class EnumerateStartTestCase(EnumerateTestCase):