SF bug #793826: using itertools.izip to mutate tuples

Avoid Armin Rigo's dastardly exercise in re-entrancy.

2 files changed, 36 insertions, 3 deletions

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 728380d..7059650 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -425,6 +425,36 @@ class TestVariousIteratorArgs(unittest.TestCase):
             self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
             self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
 
+class RegressionTests(unittest.TestCase):
+
+    def test_sf_793826(self):
+        # Fix Armin Rigo's successful efforts to wreak havoc
+
+        def mutatingtuple(tuple1, f, tuple2):
+            # this builds a tuple t which is a copy of tuple1,
+            # then calls f(t), then mutates t to be equal to tuple2
+            # (needs len(tuple1) == len(tuple2)).
+            def g(value, first=[1]):
+                if first:
+                    del first[:]
+                    f(z.next())
+                return value
+            items = list(tuple2)
+            items[1:1] = list(tuple1)
+            gen = imap(g, items)
+            z = izip(*[gen]*len(tuple1))
+            z.next()
+
+        def f(t):
+            global T
+            T = t
+            first[:] = list(T)
+
+        first = []
+        mutatingtuple((1,2,3), f, (4,5,6))
+        second = list(T)
+        self.assertEqual(first, second)
+
 
 libreftest = """ Doctest for examples in the library reference: libitertools.tex
 
@@ -559,7 +589,8 @@ False
 __test__ = {'libreftest' : libreftest}
 
 def test_main(verbose=None):
-    test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC)
+    test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
+                    RegressionTests)
     test_support.run_unittest(*test_classes)
 
     # verify reference counting
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index e63d02a..4f49fef 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -1599,16 +1599,18 @@ izip_next(izipobject *lz)
 	PyObject *item;
 
 	if (result->ob_refcnt == 1) {
+		Py_INCREF(result);
 		for (i=0 ; i < tuplesize ; i++) {
 			it = PyTuple_GET_ITEM(lz->ittuple, i);
 			assert(PyIter_Check(it));
 			item = (*it->ob_type->tp_iternext)(it);
-			if (item == NULL)
+			if (item == NULL) {
+				Py_DECREF(result);
 				return NULL;
+			}
 			Py_DECREF(PyTuple_GET_ITEM(result, i));
 			PyTuple_SET_ITEM(result, i, item);
 		}
-		Py_INCREF(result);
 	} else {
 		result = PyTuple_New(tuplesize);
 		if (result == NULL)