Commit fa220ec7 authored by HongWeipeng's avatar HongWeipeng Committed by Raymond Hettinger

Raise a RuntimeError when tee iterator is consumed from different threads (GH-15567)

parent 13f37f2b
...@@ -643,7 +643,8 @@ loops that truncate the stream. ...@@ -643,7 +643,8 @@ loops that truncate the stream.
Once :func:`tee` has made a split, the original *iterable* should not be Once :func:`tee` has made a split, the original *iterable* should not be
used anywhere else; otherwise, the *iterable* could get advanced without used anywhere else; otherwise, the *iterable* could get advanced without
the tee objects being informed. the tee objects being informed. the :func:`tee` iterator can not be consumed
from different threads, even if an underlying iterator is thread-safe.
This itertool may require significant auxiliary storage (depending on how This itertool may require significant auxiliary storage (depending on how
much temporary data needs to be stored). In general, if one iterator uses much temporary data needs to be stored). In general, if one iterator uses
......
Raise a RuntimeError when itertools.tee() iterator is consumed from different
threads. Patch by hongweipeng.
...@@ -452,6 +452,7 @@ typedef struct { ...@@ -452,6 +452,7 @@ typedef struct {
teedataobject *dataobj; teedataobject *dataobj;
int index; /* 0 <= index <= LINKCELLS */ int index; /* 0 <= index <= LINKCELLS */
PyObject *weakreflist; PyObject *weakreflist;
unsigned long thread_id;
} teeobject; } teeobject;
static PyTypeObject teedataobject_type; static PyTypeObject teedataobject_type;
...@@ -680,6 +681,11 @@ tee_next(teeobject *to) ...@@ -680,6 +681,11 @@ tee_next(teeobject *to)
{ {
PyObject *value, *link; PyObject *value, *link;
if (to->thread_id != PyThread_get_thread_ident()) {
PyErr_SetString(PyExc_RuntimeError,
"tee() iterator can not be consumed from different threads.");
return NULL;
}
if (to->index >= LINKCELLS) { if (to->index >= LINKCELLS) {
link = teedataobject_jumplink(to->dataobj); link = teedataobject_jumplink(to->dataobj);
if (link == NULL) if (link == NULL)
...@@ -713,6 +719,7 @@ tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored)) ...@@ -713,6 +719,7 @@ tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
newto->dataobj = to->dataobj; newto->dataobj = to->dataobj;
newto->index = to->index; newto->index = to->index;
newto->weakreflist = NULL; newto->weakreflist = NULL;
newto->thread_id = to->thread_id;
PyObject_GC_Track(newto); PyObject_GC_Track(newto);
return (PyObject *)newto; return (PyObject *)newto;
} }
...@@ -745,6 +752,7 @@ tee_fromiterable(PyObject *iterable) ...@@ -745,6 +752,7 @@ tee_fromiterable(PyObject *iterable)
to->index = 0; to->index = 0;
to->weakreflist = NULL; to->weakreflist = NULL;
to->thread_id = PyThread_get_thread_ident();
PyObject_GC_Track(to); PyObject_GC_Track(to);
done: done:
Py_XDECREF(it); Py_XDECREF(it);
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment