Fix line-by-line file iteration in ioutils modules

- iterator on self was yielding only the first newline
- added __next__ to match next() for py3 compat
- updated tests
This commit is contained in:
Brant Watson 2017-11-08 15:37:41 -06:00 committed by Mahmoud Hashemi
parent 115cded548
commit 55a2791b77
2 changed files with 25 additions and 10 deletions

View File

@ -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

View File

@ -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):