[pypy-commit] pypy numpypy-nditer: Make sure nditer.index does not give an index when the iterator is consumed

[rguillebert]

Author: Romain Guillebert <romain.py at gmail.com>
Branch: numpypy-nditer
Changeset: r64918:97a5ef769a6c
Date: 2013-06-17 15:43 +0200
http://bitbucket.org/pypy/pypy/changeset/97a5ef769a6c/

Log:	Make sure nditer.index does not give an index when the iterator is
	consumed

diff --git a/pypy/module/micronumpy/interp_nditer.py b/pypy/module/micronumpy/interp_nditer.py
--- a/pypy/module/micronumpy/interp_nditer.py
+++ b/pypy/module/micronumpy/interp_nditer.py
@@ -264,6 +264,7 @@
         self.reduce_ok = False
         self.zerosize_ok = False
         self.index_iter = None
+        self.done = False
         if space.isinstance_w(w_seq, space.w_tuple) or \
            space.isinstance_w(w_seq, space.w_list):
             w_seq_as_list = space.listview(w_seq)
@@ -308,6 +309,7 @@
             if not it.done():
                 break
         else:
+            self.done = True
             raise OperationError(space.w_StopIteration, space.w_None)
         res = []
         if self.index_iter:
@@ -370,6 +372,8 @@
     def descr_get_index(self, space):
         if self.tracked_index == "":
             raise OperationError(space.w_ValueError, space.wrap("Iterator does not have an index"))
+        if self.done:
+            raise OperationError(space.w_ValueError, space.wrap("Iterator is past the end"))
         return space.wrap(self.index_iter.getvalue())
 
     def descr_get_has_multi_index(self, space):
diff --git a/pypy/module/micronumpy/test/test_nditer.py b/pypy/module/micronumpy/test/test_nditer.py
--- a/pypy/module/micronumpy/test/test_nditer.py
+++ b/pypy/module/micronumpy/test/test_nditer.py
@@ -76,6 +76,12 @@
         for value in it:
             r.append((value, it.index))
         assert r == [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5)]
+        exc = None
+        try:
+            it.index
+        except ValueError, e:
+            exc = e
+        assert exc
 
         r = []
         it = nditer(a, flags=['f_index'])