diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst index 959424ff914390..76899ce753685a 100644 --- a/Doc/library/itertools.rst +++ b/Doc/library/itertools.rst @@ -628,6 +628,10 @@ loops that truncate the stream. used anywhere else; otherwise, the *iterable* could get advanced without the tee objects being informed. + ``tee`` iterators are not threadsafe. A :exc:`RuntimeError` may be + raised when using simultaneously iterators returned by the same :func:`tee` + call, even if the original *iterable* is threadsafe. + This itertool may require significant auxiliary storage (depending on how much temporary data needs to be stored). In general, if one iterator uses most or all of the data before another iterator starts, it is faster to use diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index cbbb4c4f71d3b8..721a17556f7ed7 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -11,6 +11,7 @@ from functools import reduce import sys import struct +import threading maxsize = support.MAX_Py_ssize_t minsize = -maxsize-1 @@ -1476,6 +1477,42 @@ def test_tee_del_backward(self): del forward, backward raise + def test_tee_reenter(self): + class I: + first = True + def __iter__(self): + return self + def __next__(self): + first = self.first + self.first = False + if first: + return next(b) + + a, b = tee(I()) + with self.assertRaisesRegex(RuntimeError, "tee"): + next(a) + + def test_tee_concurrent(self): + start = threading.Event() + finish = threading.Event() + class I: + def __iter__(self): + return self + def __next__(self): + start.set() + finish.wait() + + a, b = tee(I()) + thread = threading.Thread(target=next, args=[a]) + thread.start() + try: + start.wait() + with self.assertRaisesRegex(RuntimeError, "tee"): + next(b) + finally: + finish.set() + thread.join() + def test_StopIteration(self): self.assertRaises(StopIteration, next, zip()) diff --git a/Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst b/Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst new file mode 100644 index 00000000000000..64e778ee0913c2 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst @@ -0,0 +1,2 @@ +Fixed a crash in the :func:`tee` iterator when re-enter it. RuntimeError is +now raised in this case. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index bff9df608304b6..0ddd98dea12474 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -395,6 +395,7 @@ typedef struct { PyObject_HEAD PyObject *it; int numread; /* 0 <= numread <= LINKCELLS */ + int running; PyObject *nextlink; PyObject *(values[LINKCELLS]); } teedataobject; @@ -417,6 +418,7 @@ teedataobject_newinternal(PyObject *it) if (tdo == NULL) return NULL; + tdo->running = 0; tdo->numread = 0; tdo->nextlink = NULL; Py_INCREF(it); @@ -445,7 +447,14 @@ teedataobject_getitem(teedataobject *tdo, int i) else { /* this is the lead iterator, so fetch more data */ assert(i == tdo->numread); + if (tdo->running) { + PyErr_SetString(PyExc_RuntimeError, + "cannot re-enter the tee iterator"); + return NULL; + } + tdo->running = 1; value = PyIter_Next(tdo->it); + tdo->running = 0; if (value == NULL) return NULL; tdo->numread++;