# HG changeset patch
# User Arnaud Bergeron <abergeron@gmail.com>
# Date 1267481209 18000
# Node ID 0d0677773533915e90824b85688f641fbb0202cd
# Parent  e98f6b855a7f03c47e9cf0585debc99b9aa267a2
Fix bug where there would be a bunch of 0-length batches at the end under certain circumstances.

diff -r e98f6b855a7f -r 0d0677773533 datasets/dsetiter.py
--- a/datasets/dsetiter.py	Mon Mar 01 11:46:39 2010 -0500
+++ b/datasets/dsetiter.py	Mon Mar 01 17:06:49 2010 -0500
@@ -123,6 +123,7 @@
             tmpbuf = self.curfile.read(self.bufsize - len(buf))
             buf = numpy.row_stack((buf, tmpbuf))
 
+        self.cursize = len(buf)
         self.buffer = buf
         self.curpos = 0
 
@@ -144,9 +145,17 @@
             Traceback (most recent call last):
               ...
             StopIteration
-            
+            >>> d = DataIterator([DummyFile(13)], 10, 50)
+            >>> len(d.next())
+            10
+            >>> len(d.next())
+            3
+            >>> d.next()
+            Traceback (most recent call last):
+              ...
+            StopIteration
         """
-        if self.curpos >= self.bufsize:
+        if self.curpos >= self.cursize:
             self._fill_buf()
         res = self.buffer[self.curpos:self.curpos+self.batchsize]
         self.curpos += self.batchsize