From 37beb3a5c5786afbac02e390b61e362f40123eba Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Tue, 12 Mar 2019 01:08:11 +0000
Subject: [PATCH] core: teach iter_split() to break on callback returning
 False.

---
 mitogen/core.py          | 15 +++++++++++----
 tests/iter_split_test.py | 27 +++++++++++++++++++++++----
 2 files changed, 34 insertions(+), 8 deletions(-)

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()