diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index b94d688738f9e8..f0fd1d28f56f55 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -518,6 +518,15 @@ def test_count(self): self.assertEqual(next(c), -8) self.assertEqual(repr(count(10.25)), 'count(10.25)') self.assertEqual(repr(count(10.0)), 'count(10.0)') + + self.assertEqual(repr(count(maxsize)), f'count({maxsize})') + c = count(maxsize - 1) + self.assertEqual(repr(c), f'count({maxsize - 1})') + next(c) # c is now at masize + self.assertEqual(repr(c), f'count({maxsize})') + next(c) + self.assertEqual(repr(c), f'count({maxsize + 1})') + self.assertEqual(type(next(count(10.0))), float) for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5): # Test repr @@ -578,6 +587,20 @@ def test_count_with_step(self): self.assertEqual(type(next(c)), int) self.assertEqual(type(next(c)), float) + c = count(maxsize -2, 2) + self.assertEqual(repr(c), f'count({maxsize - 2}, 2)') + next(c) # c is now at masize + self.assertEqual(repr(c), f'count({maxsize}, 2)') + next(c) + self.assertEqual(repr(c), f'count({maxsize + 2}, 2)') + + c = count(maxsize + 1, -1) + self.assertEqual(repr(c), f'count({maxsize + 1}, -1)') + next(c) # c is now at masize + self.assertEqual(repr(c), f'count({maxsize}, -1)') + next(c) + self.assertEqual(repr(c), f'count({maxsize - 1}, -1)') + @threading_helper.requires_working_threading() def test_count_threading(self, step=1): # this test verifies multithreading consistency, which is diff --git a/Misc/NEWS.d/next/Library/2024-11-09-14-18-02.gh-issue-126618.Tq9FD4.rst b/Misc/NEWS.d/next/Library/2024-11-09-14-18-02.gh-issue-126618.Tq9FD4.rst new file mode 100644 index 00000000000000..da09004657b907 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2024-11-09-14-18-02.gh-issue-126618.Tq9FD4.rst @@ -0,0 +1,2 @@ +Fix the representation of :class:`itertools.count` objects when the count +value is :data:`sys.maxsize`. Patch by Bénédikt Tran. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 78fbdcdf77a923..ea3918579c0166 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3330,6 +3330,14 @@ itertools_count_impl(PyTypeObject *type, PyObject *long_cnt, assert(!fast_mode || (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1)); + if (cnt == PY_SSIZE_T_MAX && long_cnt == NULL) { + long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX); + if (long_cnt == NULL) { + Py_DECREF(long_step); + return NULL; + } + } + /* create countobject structure */ lz = (countobject *)type->tp_alloc(type, 0); if (lz == NULL) { @@ -3386,13 +3394,35 @@ count_nextlong(countobject *lz) return long_cnt; } +static inline int +count_switch_to_slow_mode(countobject *lz) +{ + assert(lz->cnt == PY_SSIZE_T_MAX); + assert(lz->long_cnt == NULL); + /* Switch to slow_mode */ + PyObject *long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX); + if (long_cnt == NULL) { + return -1; + } + lz->long_cnt = long_cnt; + return 0; +} + static PyObject * count_next(countobject *lz) { #ifndef Py_GIL_DISABLED if (lz->cnt == PY_SSIZE_T_MAX) return count_nextlong(lz); - return PyLong_FromSsize_t(lz->cnt++); + PyObject *res = PyLong_FromSsize_t(lz->cnt++); + if (lz->cnt == PY_SSIZE_T_MAX && lz->long_cnt == NULL) { + /* populate lz->long_cnt since it could be used by repr() */ + if (count_switch_to_slow_mode(lz) < 0) { + Py_DECREF(res); + return NULL; + } + } + return res; #else // free-threading version // fast mode uses compare-exchange loop @@ -3409,7 +3439,19 @@ count_next(countobject *lz) return returned; } if (_Py_atomic_compare_exchange_ssize(&lz->cnt, &cnt, cnt + 1)) { - return PyLong_FromSsize_t(cnt); + /* populate lz->long_cnt now since it could be used by repr() */ + returned = PyLong_FromSsize_t(cnt); + if (lz->cnt == PY_SSIZE_T_MAX && lz->long_cnt == NULL) { + int rc; + Py_BEGIN_CRITICAL_SECTION(lz); + rc = count_switch_to_slow_mode(lz); + Py_END_CRITICAL_SECTION(); + if (rc < 0) { + Py_DECREF(returned); + return NULL; + } + } + return returned; } } #endif