diff options
| -rw-r--r-- | Lib/test/test_itertools.py | 9 | ||||
| -rw-r--r-- | Modules/itertoolsmodule.c | 18 |
2 files changed, 26 insertions, 1 deletions
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 2f449da0bb..16e65cab65 100644 --- 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 from functools import reduce maxsize = support.MAX_Py_ssize_t minsize = -maxsize-1 @@ -352,6 +354,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, maxsize-5, maxsize+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(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)]) self.assertEqual(lzip('abc',count(start=2,step=3)), diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 3264d1f762..74396b68b8 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3049,6 +3049,22 @@ count_repr(countobject *lz) lz->long_cnt, lz->long_step); } +static PyObject * +count_reduce(countobject *lz) +{ + if (lz->cnt == PY_SSIZE_T_MAX) + return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step); + 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}, + {NULL, NULL} /* sentinel */ +}; + PyDoc_STRVAR(count_doc, "count(start=0, step=1]) --> count object\n\ \n\ @@ -3090,7 +3106,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 */ |