From a746d5e272c0fb859d2d9645d19161c58a389718 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?= <10796600+picnixz@users.noreply.github.com> Date: Sat, 9 Nov 2024 14:11:19 +0100 Subject: [PATCH 1/4] fix `itertools.count.__repr__` --- Lib/test/test_itertools.py | 12 +++++++++++ Modules/itertoolsmodule.c | 42 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index a52e1d3fa142d9..f12e4cebff6ab2 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -516,6 +516,13 @@ 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) # now c is at maxsize + self.assertEqual(repr(c), f'count({maxsize})') + 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 @@ -570,6 +577,11 @@ def test_count_with_step(self): self.assertEqual(type(next(c)), int) self.assertEqual(type(next(c)), float) + c = count(maxsize + 1, -1) + self.assertEqual(repr(c), f'count({maxsize + 1}, -1)') + next(c) # now c is at maxsize + self.assertEqual(repr(c), f'count({maxsize}, -1)') + @threading_helper.requires_working_threading() def test_count_threading(self, step=1): # this test verifies multithreading consistency, which is diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 1201fa094902d7..3e1f68f67255ed 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3327,6 +3327,10 @@ 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); + } + /* create countobject structure */ lz = (countobject *)type->tp_alloc(type, 0); if (lz == NULL) { @@ -3383,13 +3387,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 @@ -3406,7 +3432,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 From 251ae19a1db6162d580f7ea7f8f4c80ddd210fa3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?= <10796600+picnixz@users.noreply.github.com> Date: Sat, 9 Nov 2024 14:18:06 +0100 Subject: [PATCH 2/4] blurb --- .../next/Library/2024-11-09-14-18-02.gh-issue-126618.Tq9FD4.rst | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 Misc/NEWS.d/next/Library/2024-11-09-14-18-02.gh-issue-126618.Tq9FD4.rst 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. From 83ebc7c35fcd639f452d34e67223a9422cfd4b5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?= <10796600+picnixz@users.noreply.github.com> Date: Sat, 9 Nov 2024 14:43:04 +0100 Subject: [PATCH 3/4] correctly propagate exception --- Modules/itertoolsmodule.c | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 3e1f68f67255ed..23527761d0020d 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3329,6 +3329,10 @@ itertools_count_impl(PyTypeObject *type, PyObject *long_cnt, 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 */ From 2346a30abef7606589add9c56ae346dfd82aca11 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?= <10796600+picnixz@users.noreply.github.com> Date: Tue, 12 Nov 2024 14:28:14 +0100 Subject: [PATCH 4/4] improve test coverage --- Lib/test/test_itertools.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 8a883beec454f2..f0fd1d28f56f55 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -522,8 +522,10 @@ def test_count(self): self.assertEqual(repr(count(maxsize)), f'count({maxsize})') c = count(maxsize - 1) self.assertEqual(repr(c), f'count({maxsize - 1})') - next(c) # now c is at maxsize + 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): @@ -585,10 +587,19 @@ 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) # now c is at maxsize + 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):