diff --git a/mitogen/core.py b/mitogen/core.py index 0890735d..a715158b 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -634,16 +634,23 @@ def iter_split(buf, delim, func): """ Invoke `func(s)` for each `delim`-delimited chunk in the potentially large `buf`, avoiding intermediate lists and quadratic string operations. Return - the trailing undelimited portion of `buf`. + the trailing undelimited portion of `buf`, or any unprocessed portion of + `buf` after `func(s)` returned :data:`False`. + + :returns: + `(trailer, cont)`, where `cont` is :data:`False` if the last call to + `func(s)` returned :data:`False`. """ dlen = len(delim) start = 0 - while True: + cont = True + while cont: nl = buf.find(delim, start) if nl == -1: - return buf[start:] - func(buf[start:nl]) + break + cont = not func(buf[start:nl]) is False start = nl + dlen + return buf[start:], cont class Py24Pickler(py_pickle.Pickler): diff --git a/tests/iter_split_test.py b/tests/iter_split_test.py index a385b2f1..6178ef0c 100644 --- a/tests/iter_split_test.py +++ b/tests/iter_split_test.py @@ -12,30 +12,49 @@ class IterSplitTest(unittest2.TestCase): def test_empty_buffer(self): lst = [] - trailer = self.func(buf='', delim='\n', func=lst.append) + trailer, cont = self.func(buf='', delim='\n', func=lst.append) + self.assertTrue(cont) self.assertEquals('', trailer) self.assertEquals([], lst) def test_empty_line(self): lst = [] - trailer = self.func(buf='\n', delim='\n', func=lst.append) + trailer, cont = self.func(buf='\n', delim='\n', func=lst.append) + self.assertTrue(cont) self.assertEquals('', trailer) self.assertEquals([''], lst) def test_one_line(self): buf = 'xxxx\n' lst = [] - trailer = self.func(buf=buf, delim='\n', func=lst.append) + trailer, cont = self.func(buf=buf, delim='\n', func=lst.append) + self.assertTrue(cont) self.assertEquals('', trailer) self.assertEquals(lst, ['xxxx']) def test_one_incomplete(self): buf = 'xxxx\nyy' lst = [] - trailer = self.func(buf=buf, delim='\n', func=lst.append) + trailer, cont = self.func(buf=buf, delim='\n', func=lst.append) + self.assertTrue(cont) self.assertEquals('yy', trailer) self.assertEquals(lst, ['xxxx']) + def test_returns_false_immediately(self): + buf = 'xxxx\nyy' + func = lambda buf: False + trailer, cont = self.func(buf=buf, delim='\n', func=func) + self.assertFalse(cont) + self.assertEquals('yy', trailer) + + def test_returns_false_second_call(self): + buf = 'xxxx\nyy\nzz' + it = iter([True, False]) + func = lambda buf: next(it) + trailer, cont = self.func(buf=buf, delim='\n', func=func) + self.assertFalse(cont) + self.assertEquals('zz', trailer) + if __name__ == '__main__': unittest2.main()