From 35722a93768e942be31270fe60acff138d878b26 Mon Sep 17 00:00:00 2001 From: Kristjan Valur Jonsson Date: Wed, 30 Mar 2011 11:04:28 +0000 Subject: Bugfix: Properly test for errors from PyLong_AsLong() in itertools.cycle. ti can raise an exception even if PyLong_Check() has succeeded. --- Lib/test/test_itertools.py | 5 +++++ Modules/itertoolsmodule.c | 9 ++++++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index da101e1..9ce54d0 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -339,6 +339,8 @@ class TestBasicOps(unittest.TestCase): list(range(maxsize-5, maxsize+5))) self.assertEqual(list(islice(count(-maxsize-5), 10)), list(range(-maxsize-5, -maxsize+5))) + self.assertEqual(list(islice(count(10, maxsize+5), 3)), + list(range(10, 10+3*(maxsize+5), maxsize+5))) c = count(3) self.assertEqual(repr(c), 'count(3)') next(c) @@ -361,6 +363,9 @@ class TestBasicOps(unittest.TestCase): self.assertEqual(next(copy.deepcopy(c)), value) self.assertEqual(next(pickle.loads(pickle.dumps(c))), value) + #check proper internal error handling for large "step' sizes + count(1, maxsize+5); sys.exc_info() + 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 d5336f2..e7a7fbe 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -2918,6 +2918,7 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) Py_ssize_t cnt = 0; PyObject *long_cnt = NULL; PyObject *long_step = NULL; + long step; static char *kwlist[] = {"start", "step", 0}; if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count", @@ -2955,9 +2956,11 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) assert(long_cnt != NULL && long_step != NULL); /* Fast mode only works when the step is 1 */ - if (!PyLong_Check(long_step) || - PyLong_AS_LONG(long_step) != 1) { - slow_mode = 1; + step = PyLong_AsLong(long_step); + if (step != 1) { + slow_mode = 1; + if (step == -1 && PyErr_Occurred()) + PyErr_Clear(); } if (slow_mode) -- cgit v0.12