diff datasets/dsetiter.py @ 189:0d0677773533

Fix bug where there would be a bunch of 0-length batches at the end under certain circumstances.
author Arnaud Bergeron <abergeron@gmail.com>
date Mon, 01 Mar 2010 17:06:49 -0500
parents 76bc047df5ee
children 1adfafdc3d57
line wrap: on
line diff
--- 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