gh-126618: Fix representation of itertools.count(sys.maxsize)

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

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.

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