[pypy-svn] r56446 - in pypy/dist/pypy/module/itertools: . test

[adurdin at codespeak.net]

Author: adurdin
Date: Fri Jul 11 14:54:36 2008
New Revision: 56446

Modified:
   pypy/dist/pypy/module/itertools/interp_itertools.py
   pypy/dist/pypy/module/itertools/test/test_itertools.py
Log:
(adurdin, jlg) interp_itertools.izip behaves per spec after exhausting.


Modified: pypy/dist/pypy/module/itertools/interp_itertools.py
==============================================================================
--- pypy/dist/pypy/module/itertools/interp_itertools.py	(original)
+++ pypy/dist/pypy/module/itertools/interp_itertools.py	Fri Jul 11 14:54:36 2008
@@ -465,7 +465,12 @@
         if not self.iterators_w:
             raise OperationError(self.space.w_StopIteration, self.space.w_None)
 
-        w_objects = [self.space.next(w_it) for w_it in self.iterators_w]
+        try:
+            w_objects = [self.space.next(w_it) for w_it in self.iterators_w]
+        except OperationError, e:
+            if e.match(self.space, self.space.w_StopIteration):
+                self.iterators_w = None
+            raise
 
         if self.identity_fun:
             return self.space.newtuple(w_objects)

Modified: pypy/dist/pypy/module/itertools/test/test_itertools.py
==============================================================================
--- pypy/dist/pypy/module/itertools/test/test_itertools.py	(original)
+++ pypy/dist/pypy/module/itertools/test/test_itertools.py	Fri Jul 11 14:54:36 2008
@@ -327,7 +327,17 @@
         
         it = itertools.izip([], [], [1], [])
         raises(StopIteration, it.next)
-        
+
+        # Up to one additional item may be consumed per iterable, as per python docs
+        it1 = iter([1, 2, 3, 4, 5, 6])
+        it2 = iter([5, 6])
+        it = itertools.izip(it1, it2)
+        for x in [(1, 5), (2, 6)]:
+            assert it.next() == x
+        raises(StopIteration, it.next)
+        assert it1.next() == 4
+        raises(StopIteration, it.next)
+        assert it1.next() == 5
 
     def test_docstrings(self):
         import itertools