gh-123471: Make itertools.chain thread-safe (#135689)

This commit is contained in:
Pieter Eendebak 2025-06-30 13:06:58 +02:00 committed by GitHub
parent 536a5ff153
commit 0533c1faf2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 43 additions and 4 deletions

View file

@ -1,6 +1,6 @@
import unittest
from threading import Thread, Barrier
from itertools import batched, cycle
from itertools import batched, chain, cycle
from test.support import threading_helper
@ -17,7 +17,7 @@ class ItertoolsThreading(unittest.TestCase):
barrier.wait()
while True:
try:
_ = next(it)
next(it)
except StopIteration:
break
@ -62,6 +62,34 @@ class ItertoolsThreading(unittest.TestCase):
barrier.reset()
@threading_helper.reap_threads
def test_chain(self):
number_of_threads = 6
number_of_iterations = 20
barrier = Barrier(number_of_threads)
def work(it):
barrier.wait()
while True:
try:
next(it)
except StopIteration:
break
data = [(1, )] * 200
for it in range(number_of_iterations):
chain_iterator = chain(*data)
worker_threads = []
for ii in range(number_of_threads):
worker_threads.append(
Thread(target=work, args=[chain_iterator]))
with threading_helper.start_threads(worker_threads):
pass
barrier.reset()
if __name__ == "__main__":
unittest.main()

View file

@ -0,0 +1 @@
Make concurrent iterations over :class:`itertools.chain` safe under :term:`free threading`.

View file

@ -1880,8 +1880,8 @@ chain_traverse(PyObject *op, visitproc visit, void *arg)
return 0;
}
static PyObject *
chain_next(PyObject *op)
static inline PyObject *
chain_next_lock_held(PyObject *op)
{
chainobject *lz = chainobject_CAST(op);
PyObject *item;
@ -1919,6 +1919,16 @@ chain_next(PyObject *op)
return NULL;
}
static PyObject *
chain_next(PyObject *op)
{
PyObject *result;
Py_BEGIN_CRITICAL_SECTION(op);
result = chain_next_lock_held(op);
Py_END_CRITICAL_SECTION()
return result;
}
PyDoc_STRVAR(chain_doc,
"chain(*iterables)\n\
--\n\