Skip to content

Commit 847d1c2

Browse files
gh-123471: Make itertools.product and itertools.combinations thread-safe (#132814)
Co-authored-by: Kumar Aditya <[email protected]>
1 parent b1056c2 commit 847d1c2

File tree

3 files changed

+74
-2
lines changed

3 files changed

+74
-2
lines changed
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
import unittest
2+
from threading import Thread, Barrier
3+
from itertools import combinations, product
4+
from test.support import threading_helper
5+
6+
7+
threading_helper.requires_working_threading(module=True)
8+
9+
def test_concurrent_iteration(iterator, number_of_threads):
10+
barrier = Barrier(number_of_threads)
11+
def iterator_worker(it):
12+
barrier.wait()
13+
while True:
14+
try:
15+
_ = next(it)
16+
except StopIteration:
17+
return
18+
19+
worker_threads = []
20+
for ii in range(number_of_threads):
21+
worker_threads.append(
22+
Thread(target=iterator_worker, args=[iterator]))
23+
24+
with threading_helper.start_threads(worker_threads):
25+
pass
26+
27+
barrier.reset()
28+
29+
class ItertoolsThreading(unittest.TestCase):
30+
31+
@threading_helper.reap_threads
32+
def test_combinations(self):
33+
number_of_threads = 10
34+
number_of_iterations = 24
35+
36+
for it in range(number_of_iterations):
37+
iterator = combinations((1, 2, 3, 4, 5), 2)
38+
test_concurrent_iteration(iterator, number_of_threads)
39+
40+
@threading_helper.reap_threads
41+
def test_product(self):
42+
number_of_threads = 10
43+
number_of_iterations = 24
44+
45+
for it in range(number_of_iterations):
46+
iterator = product((1, 2, 3, 4, 5), (10, 20, 30))
47+
test_concurrent_iteration(iterator, number_of_threads)
48+
49+
50+
if __name__ == "__main__":
51+
unittest.main()
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Make concurrent iterations over :class:`itertools.combinations` and :class:`itertools.product` safe under free-threading.

Modules/itertoolsmodule.c

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2096,7 +2096,7 @@ product_traverse(PyObject *op, visitproc visit, void *arg)
20962096
}
20972097

20982098
static PyObject *
2099-
product_next(PyObject *op)
2099+
product_next_lock_held(PyObject *op)
21002100
{
21012101
productobject *lz = productobject_CAST(op);
21022102
PyObject *pool;
@@ -2182,6 +2182,16 @@ product_next(PyObject *op)
21822182
return NULL;
21832183
}
21842184

2185+
static PyObject *
2186+
product_next(PyObject *op)
2187+
{
2188+
PyObject *result;
2189+
Py_BEGIN_CRITICAL_SECTION(op);
2190+
result = product_next_lock_held(op);
2191+
Py_END_CRITICAL_SECTION()
2192+
return result;
2193+
}
2194+
21852195
static PyMethodDef product_methods[] = {
21862196
{"__sizeof__", product_sizeof, METH_NOARGS, sizeof_doc},
21872197
{NULL, NULL} /* sentinel */
@@ -2329,7 +2339,7 @@ combinations_traverse(PyObject *op, visitproc visit, void *arg)
23292339
}
23302340

23312341
static PyObject *
2332-
combinations_next(PyObject *op)
2342+
combinations_next_lock_held(PyObject *op)
23332343
{
23342344
combinationsobject *co = combinationsobject_CAST(op);
23352345
PyObject *elem;
@@ -2414,6 +2424,16 @@ combinations_next(PyObject *op)
24142424
return NULL;
24152425
}
24162426

2427+
static PyObject *
2428+
combinations_next(PyObject *op)
2429+
{
2430+
PyObject *result;
2431+
Py_BEGIN_CRITICAL_SECTION(op);
2432+
result = combinations_next_lock_held(op);
2433+
Py_END_CRITICAL_SECTION()
2434+
return result;
2435+
}
2436+
24172437
static PyMethodDef combinations_methods[] = {
24182438
{"__sizeof__", combinations_sizeof, METH_NOARGS, sizeof_doc},
24192439
{NULL, NULL} /* sentinel */

0 commit comments

Comments
 (0)