Skip to content

Commit

Permalink
bpo-34410: Fix a crash in the tee iterator when re-enter it. (pythonG…
Browse files Browse the repository at this point in the history
…H-15625)

RuntimeError is now raised in this case.
  • Loading branch information
serhiy-storchaka authored and lisroach committed Sep 9, 2019
1 parent c528ba4 commit 988fb1c
Show file tree
Hide file tree
Showing 4 changed files with 52 additions and 0 deletions.
4 changes: 4 additions & 0 deletions Doc/library/itertools.rst
Original file line number Diff line number Diff line change
Expand Up @@ -645,6 +645,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
Expand Down
37 changes: 37 additions & 0 deletions Lib/test/test_itertools.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from functools import reduce
import sys
import struct
import threading
maxsize = support.MAX_Py_ssize_t
minsize = -maxsize-1

Expand Down Expand Up @@ -1494,6 +1495,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())

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fixed a crash in the :func:`tee` iterator when re-enter it. RuntimeError is
now raised in this case.
9 changes: 9 additions & 0 deletions Modules/itertoolsmodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -443,6 +443,7 @@ typedef struct {
PyObject_HEAD
PyObject *it;
int numread; /* 0 <= numread <= LINKCELLS */
int running;
PyObject *nextlink;
PyObject *(values[LINKCELLS]);
} teedataobject;
Expand All @@ -465,6 +466,7 @@ teedataobject_newinternal(PyObject *it)
if (tdo == NULL)
return NULL;

tdo->running = 0;
tdo->numread = 0;
tdo->nextlink = NULL;
Py_INCREF(it);
Expand Down Expand Up @@ -493,7 +495,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++;
Expand Down

0 comments on commit 988fb1c

Please sign in to comment.