Issue 7410: deepcopy of itertools.count resets the count

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index f8bdb8bdec9fe713d0c28c5ac9a2aaca479c186b..03f1d85111ef27a1bbdb00396774824b4b9416d5 100644 (file)
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -7,6 +7,8 @@ from fractions import Fraction
 import sys
 import operator
 import random
+import copy
+import pickle
 maxsize = test_support.MAX_Py_ssize_t
 minsize = -maxsize-1
 
@@ -346,6 +348,13 @@ class TestBasicOps(unittest.TestCase):
             r2 = 'count(%r)'.__mod__(i).replace('L', '')
             self.assertEqual(r1, r2)
 
+        # check copy, deepcopy, pickle
+        for value in -3, 3, sys.maxint-5, sys.maxint+5:
+            c = count(value)
+            self.assertEqual(next(copy.copy(c)), value)
+            self.assertEqual(next(copy.deepcopy(c)), value)
+            self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
+
     def test_count_with_stride(self):
         self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
         self.assertEqual(zip('abc',count(start=2,step=3)),

diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 1a0d44dfc7b53e44eb2229a3a5ea29e9afa706df..6b5c249f0056cba98a68b90d52082230fb5bcc54 100644 (file)
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -3375,6 +3375,21 @@ count_repr(countobject *lz)
        return result;
 }
 
+static PyObject *
+count_reduce(countobject *lz)
+{
+       if (lz->cnt == PY_SSIZE_T_MAX)
+               return Py_BuildValue("O(O)", Py_TYPE(lz), lz->long_cnt);
+       return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
+}
+
+PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
+
+static PyMethodDef count_methods[] = {
+       {"__reduce__",  (PyCFunction)count_reduce,      METH_NOARGS,
+        count_reduce_doc},
+};
+
 PyDoc_STRVAR(count_doc,
                         "count(start=0, step=1]) --> count object\n\
 \n\
@@ -3416,7 +3431,7 @@ static PyTypeObject count_type = {
        0,                              /* tp_weaklistoffset */
        PyObject_SelfIter,              /* tp_iter */
        (iternextfunc)count_next,       /* tp_iternext */
-       0,                              /* tp_methods */
+       count_methods,                          /* tp_methods */
        0,                              /* tp_members */
        0,                              /* tp_getset */
        0,                              /* tp_base */