comparison 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
comparison
equal deleted inserted replaced
188:e98f6b855a7f 189:0d0677773533
121 raise 121 raise
122 break 122 break
123 tmpbuf = self.curfile.read(self.bufsize - len(buf)) 123 tmpbuf = self.curfile.read(self.bufsize - len(buf))
124 buf = numpy.row_stack((buf, tmpbuf)) 124 buf = numpy.row_stack((buf, tmpbuf))
125 125
126 self.cursize = len(buf)
126 self.buffer = buf 127 self.buffer = buf
127 self.curpos = 0 128 self.curpos = 0
128 129
129 def __next__(self): 130 def __next__(self):
130 r""" 131 r"""
142 StopIteration 143 StopIteration
143 >>> d.next() 144 >>> d.next()
144 Traceback (most recent call last): 145 Traceback (most recent call last):
145 ... 146 ...
146 StopIteration 147 StopIteration
147 148 >>> d = DataIterator([DummyFile(13)], 10, 50)
149 >>> len(d.next())
150 10
151 >>> len(d.next())
152 3
153 >>> d.next()
154 Traceback (most recent call last):
155 ...
156 StopIteration
148 """ 157 """
149 if self.curpos >= self.bufsize: 158 if self.curpos >= self.cursize:
150 self._fill_buf() 159 self._fill_buf()
151 res = self.buffer[self.curpos:self.curpos+self.batchsize] 160 res = self.buffer[self.curpos:self.curpos+self.batchsize]
152 self.curpos += self.batchsize 161 self.curpos += self.batchsize
153 return res 162 return res
154 163