gh-124397: Add threading.iter_locked

Doc/library/threading.rst

@@ -208,6 +208,48 @@ This module defines the following functions:
    of the result, even when terminated.
 
 
+.. function:: iter_locked(iterable)
+
+   Convert an iterable into an iterator that performs iteration using locks.
+
+   The ``iter_locked`` makes non-atomic iterators atomic::
+
+       class non_atomic_iterator:
+
+            def __init__(self, it):
+                self.it = iter(it)
+
+            def __iter__(self):
+                return self
+
+            def __next__(self):
+                a = next(self.it)
+                b = next(self.it)
+                return a, b
+
+        atomic_iterator = iter_locked(non_atomic_iterator())
+
+   The ``iter_locked`` allows concurrent iteration over generator objects. For example::
+
+        def count():
+            i = 0
+            while True:
+                i += 1
+                yield i
+        concurrent_iterator = iter_locked(count())
+
+   The implementation is roughly equivalent to::
+
+        class iter_locked(Iterator):
+            def __init__(self, it):
+                self._it = iter(it)
+                self._lock = Lock()
+            def __next__(self):
+                with self._lock:
+                    return next(self._it)
+
+   .. versionadded:: next
+
 .. function:: main_thread()
 
    Return the main :class:`Thread` object.  In normal conditions, the

Lib/test/test_free_threading/test_threading_iter_locked.py

@@ -0,0 +1,85 @@
+import time
+import unittest
+from threading import Thread, Barrier, iter_locked
+from test.support import threading_helper
+
+
+threading_helper.requires_working_threading(module=True)
+
+class non_atomic_iterator:
+
+    def __init__(self, it):
+        self.it = iter(it)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        a = next(self.it)
+        t = time.perf_counter() + 1e-6
+        while time.perf_counter() < t:
+            pass
+        b = next(self.it)
+        return a, b
+
+def count():
+    i = 0
+    while True:
+        i += 1
+        yield i
+
+class iter_lockedThreading(unittest.TestCase):
+
+    @threading_helper.reap_threads
+    def test_iter_locked(self):
+        number_of_threads = 10
+        number_of_iterations = 8
+        barrier = Barrier(number_of_threads)
+        def work(it):
+            while True:
+                try:
+                    a, b = next(it)
+                    assert a + 1 == b
+                except StopIteration:
+                    break
+
+        data = tuple(range(400))
+        for it in range(number_of_iterations):
+            iter_locked_iterator = iter_locked(non_atomic_iterator(data,))
+            worker_threads = []
+            for ii in range(number_of_threads):
+                worker_threads.append(
+                    Thread(target=work, args=[iter_locked_iterator]))
+
+            with threading_helper.start_threads(worker_threads):
+                pass
+
+            barrier.reset()
+
+    @threading_helper.reap_threads
+    def test_iter_locked_generator(self):
+        number_of_threads = 5
+        number_of_iterations = 4
+        barrier = Barrier(number_of_threads)
+        def work(it):
+            barrier.wait()
+            for _ in range(1_000):
+                try:
+                    next(it)
+                except StopIteration:
+                    break
+
+        for it in range(number_of_iterations):
+            generator = iter_locked(count())
+            worker_threads = []
+            for ii in range(number_of_threads):
+                worker_threads.append(
+                    Thread(target=work, args=[generator]))
+
+            with threading_helper.start_threads(worker_threads):
+                pass
+
+            barrier.reset()
+
+if __name__ == "__main__":
+    unittest.main()

Lib/test/test_threading.py

@@ -2496,6 +2496,45 @@ def run_last():
         self.assertIn("RuntimeError: can't register atexit after shutdown",
                 err.decode())
 
+class IterLockedTests(unittest.TestCase):
+
+     def test_iter_locked(self):
+            for s in ("123", [], [1, 2, 3], tuple(), (1, 2, 3)):
+                expected = list(s)
+                actual = list(threading.iter_locked(s))
+                self.assertEqual(actual, expected)
+            for arg in [1, None, True, sys]:
+                self.assertRaises(TypeError, threading.iter_locked, arg)
+
+     def test_iter_locked_recursive_generator(self):
+            def g():
+                yield next(it)
+
+            it = threading.iter_locked(g())
+            with self.assertRaises(ValueError):
+                next(it)
+
+     def test_iter_locked_recursive_iterator(self):
+        class evil_it:
+            def __init__(self):
+                self.it = iter(range(100))
+                self.it2 = None
+
+            def __iter__(self):
+                return self
+
+            def __next__(self):
+                if self.it2 is not None:
+                    it2, self.it2 = self.it2, None
+                    return next(it2), next(it2)
+                return next(self.it)
+
+        a = evil_it()
+        il = threading.iter_locked(a)
+        assert next(il) == 0
+        a.it2 = il
+        assert next(il) == (1, 2)
+
 
 if __name__ == "__main__":
     unittest.main()

Lib/threading.py

@@ -29,7 +29,7 @@
            'Barrier', 'BrokenBarrierError', 'Timer', 'ThreadError',
            'setprofile', 'settrace', 'local', 'stack_size',
            'excepthook', 'ExceptHookArgs', 'gettrace', 'getprofile',
-           'setprofile_all_threads','settrace_all_threads']
+           'setprofile_all_threads','settrace_all_threads', 'iter_locked']
 
 # Rename some stuff so "from threading import *" is safe
 _start_joinable_thread = _thread.start_joinable_thread
@@ -42,6 +42,7 @@
 get_ident = _thread.get_ident
 _get_main_thread_ident = _thread._get_main_thread_ident
 _is_main_interpreter = _thread._is_main_interpreter
+iter_locked = _thread.iter_locked
 try:
     get_native_id = _thread.get_native_id
     _HAVE_THREAD_NATIVE_ID = True

Misc/NEWS.d/next/Library/2025-05-11-19-54-05.gh-issue-124397.MwtP6-.rst

@@ -0,0 +1 @@
+Add :meth:`threading.iter_locked` to make concurrent iteration over an iterable execute using a lock.

Modules/_threadmodule.c

@@ -27,6 +27,7 @@ static struct PyModuleDef thread_module;
 // Module state
 typedef struct {
     PyTypeObject *excepthook_type;
+    PyTypeObject *iter_locked_type;
     PyTypeObject *lock_type;
     PyTypeObject *rlock_type;
     PyTypeObject *local_type;
@@ -68,10 +69,12 @@ static PF_SET_THREAD_DESCRIPTION pSetThreadDescription = NULL;
 
 /*[clinic input]
 module _thread
+class _thread.iter_locked "iter_locked_object *" "clinic_state()->iter_locked_type"
 class _thread.lock "lockobject *" "clinic_state()->lock_type"
 class _thread.RLock "rlockobject *" "clinic_state()->rlock_type"
+
 [clinic start generated code]*/
-/*[clinic end generated code: output=da39a3ee5e6b4b0d input=c5a0f8c492a0c263]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=1a4bf65233f83eae]*/
 
 #define clinic_state() get_thread_state_by_cls(type)
 #include "clinic/_threadmodule.c.h"
@@ -751,6 +754,102 @@ static PyType_Spec ThreadHandle_Type_spec = {
     ThreadHandle_Type_slots,
 };
 
+/* iter_locked object **************************************************************/
+
+typedef struct {
+    PyObject_HEAD
+    PyObject *it;
+    _PyRecursiveMutex lock;
+} iter_locked_object;
+
+#define iter_locked_object_CAST(op)    ((iter_locked_object *)(op))
+
+/*[clinic input]
+@classmethod
+_thread.iter_locked.__new__
+    iterable: object
+    /
+Make an iterator thread-safe.
+[clinic start generated code]*/
+
+static PyObject *
+_thread_iter_locked_impl(PyTypeObject *type, PyObject *iterable)
+/*[clinic end generated code: output=4a8ad5a25f7c09ba input=ae6124177726e809]*/
+{
+    /* Get iterator. */
+    PyObject *it = PyObject_GetIter(iterable);
+    if (it == NULL)
+        return NULL;
+
+    iter_locked_object *il = (iter_locked_object *)type->tp_alloc(type, 0);
+    if (il == NULL) {
+        Py_DECREF(it);
+        return NULL;
+    }
+    il->it = it;
+    il->lock = (_PyRecursiveMutex){0};
+
+    return (PyObject *)il;
+}
+
+static void
+iter_locked_dealloc(PyObject *op)
+{
+    iter_locked_object *il = iter_locked_object_CAST(op);
+    PyTypeObject *tp = Py_TYPE(il);
+    PyObject_GC_UnTrack(il);
+    Py_DECREF(il->it);
+    tp->tp_free(il);
+    Py_DECREF(tp);
+}
+
+static int
+iter_locked_traverse(PyObject *op, visitproc visit, void *arg)
+{
+    iter_locked_object *lz = iter_locked_object_CAST(op);
+    Py_VISIT(Py_TYPE(lz));
+    Py_VISIT(lz->it);
+    return 0;
+}
+
+static PyObject *
+iter_locked_next(PyObject *op)
+{
+    iter_locked_object *lz = iter_locked_object_CAST(op);
+    PyObject *result = NULL;
+
+    _PyRecursiveMutex_Lock(&(lz->lock));
+
+    int v = PyIter_NextItem(lz->it, &result);
+    if (v == -1) {
+        /* Note:  StopIteration is already cleared by PyIter_Next() */
+        /* If PyErr_Occurred() we will also return NULL*/
+    }
+    _PyRecursiveMutex_Unlock(&(lz->lock));
+
+    return result;
+}
+
+static PyType_Slot iter_locked_slots[] = {
+    {Py_tp_dealloc, iter_locked_dealloc},
+    {Py_tp_getattro, PyObject_GenericGetAttr},
+    {Py_tp_doc, (void *)_thread_iter_locked__doc__},
+    {Py_tp_traverse, iter_locked_traverse},
+    {Py_tp_iter, PyObject_SelfIter},
+    {Py_tp_iternext, iter_locked_next},
+    {Py_tp_new, _thread_iter_locked},
+    {Py_tp_free, PyObject_GC_Del},
+    {0, NULL},
+};
+
+static PyType_Spec iter_locked_spec = {
+    .name = "threading.iter_locked",
+    .basicsize = sizeof(iter_locked_object),
+    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
+              Py_TPFLAGS_IMMUTABLETYPE),
+    .slots = iter_locked_slots,
+};
+
 /* Lock objects */
 
 typedef struct {
@@ -2660,6 +2759,15 @@ thread_module_exec(PyObject *module)
         return -1;
     }
 
+    // iter_locked
+    state->iter_locked_type = (PyTypeObject *)PyType_FromModuleAndSpec(module, &iter_locked_spec, NULL);
+    if (state->iter_locked_type == NULL) {
+        return -1;
+    }
+    if (PyModule_AddType(module, state->iter_locked_type) < 0) {
+        return -1;
+    }
+
     // Lock
     state->lock_type = (PyTypeObject *)PyType_FromModuleAndSpec(module, &lock_type_spec, NULL);
     if (state->lock_type == NULL) {
@@ -2766,6 +2874,7 @@ thread_module_traverse(PyObject *module, visitproc visit, void *arg)
 {
     thread_module_state *state = get_thread_state(module);
     Py_VISIT(state->excepthook_type);
+    Py_VISIT(state->iter_locked_type);
     Py_VISIT(state->lock_type);
     Py_VISIT(state->rlock_type);
     Py_VISIT(state->local_type);
@@ -2779,6 +2888,7 @@ thread_module_clear(PyObject *module)
 {
     thread_module_state *state = get_thread_state(module);
     Py_CLEAR(state->excepthook_type);
+    Py_CLEAR(state->iter_locked_type);
     Py_CLEAR(state->lock_type);
     Py_CLEAR(state->rlock_type);
     Py_CLEAR(state->local_type);