diff --git a/boltons/ioutils.py b/boltons/ioutils.py index 47f2b6a..22e2444 100644 --- a/boltons/ioutils.py +++ b/boltons/ioutils.py @@ -116,7 +116,15 @@ class SpooledIOBase(object): return self.buffer.isatty() def next(self): - return self.readline() + line = self.readline() + if not line: + pos = self.buffer.tell() + self.buffer.seek(0, os.SEEK_END) + if pos == self.buffer.tell(): + raise StopIteration + else: + self.buffer.seek(pos) + return line @property def closed(self): @@ -161,11 +169,13 @@ class SpooledIOBase(object): self.seek(pos) return val + __next__ = next + def __len__(self): return self.len def __iter__(self): - yield self.readline() + return self def __enter__(self): return self @@ -184,7 +194,7 @@ class SpooledIOBase(object): def __bool__(self): return True - __nonzero__=__bool__ + __nonzero__ = __bool__ class SpooledBytesIO(SpooledIOBase): diff --git a/tests/test_ioutils.py b/tests/test_ioutils.py index a798dcb..98f007e 100644 --- a/tests/test_ioutils.py +++ b/tests/test_ioutils.py @@ -142,13 +142,6 @@ class BaseTestMixin(object): self.spooled_flo.softspace = True self.assertTrue(self.spooled_flo.softspace) - def test_iter(self): - """Make sure iter works as expected""" - self.spooled_flo.write(self.test_str) - self.spooled_flo.seek(0) - self.assertEqual([x for x in self.spooled_flo][0], self.test_str) - self.assertEqual([x for x in self.spooled_flo][0], self.data_type()) - def test_buf_property(self): """'buf' property returns the same value as getvalue()""" self.assertEqual(self.spooled_flo.buf, self.spooled_flo.getvalue()) @@ -273,6 +266,12 @@ class TestSpooledBytesIO(TestCase, BaseTestMixin, AssertionsMixin): self.assertTrue('content.txt' in doc.namelist()) doc.close() + def test_iter(self): + """Make sure iter works as expected""" + self.spooled_flo.write(b"a\nb") + self.spooled_flo.seek(0) + self.assertEqual([x for x in self.spooled_flo], [b"a\n", b"b"]) + class TestSpooledStringIO(TestCase, BaseTestMixin, AssertionsMixin): linesep = os.linesep @@ -405,6 +404,12 @@ class TestSpooledStringIO(TestCase, BaseTestMixin, AssertionsMixin): self.spooled_flo.seek(0) self.assertEqual(self.spooled_flo.read(3), test_str) + def test_iter(self): + """Make sure iter works as expected""" + self.spooled_flo.write(u"a\nb") + self.spooled_flo.seek(0) + self.assertEqual([x for x in self.spooled_flo], [u"a\n", u"b"]) + class TestMultiFileReader(TestCase): def test_read_seek_bytes(self):