util: Add helper for splitting up huge iterables

This is used to split a large iterable into more manageable chunks.
This is important when a consumer can only handle so much input at once

e.g. A program can take a large, but finite number of arguments, xargs
exists to handle this for shell programs.

2 files changed, 21 insertions, 0 deletions

diff --git a/morphlib/util.py b/morphlib/util.py
index 04df0633..dd2d05e1 100644
--- a/morphlib/util.py
+++ b/morphlib/util.py
@@ -13,6 +13,7 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
+import itertools
 import os
 import re
 
@@ -412,3 +413,12 @@ def get_host_architecture(): # pragma: no cover
 def sanitize_environment(env):
     for k in env:
         env[k] = str(env[k])
+
+def iter_trickle(iterable, limit):
+    '''Split an iterable up into `limit` length chunks.'''
+    it = iter(iterable)
+    while True:
+        buf = list(itertools.islice(it, limit))
+        if len(buf) == 0:
+           break
+        yield buf
diff --git a/morphlib/util_tests.py b/morphlib/util_tests.py
index 2ad9e8aa..fbf7f27b 100644
--- a/morphlib/util_tests.py
+++ b/morphlib/util_tests.py
@@ -110,3 +110,14 @@ class ParseEnvironmentPairsTests(unittest.TestCase):
         d = { 'a': 1 }
         morphlib.util.sanitize_environment(d)
         self.assertTrue(isinstance(d['a'], str))
+
+class IterTrickleTests(unittest.TestCase):
+
+    def test_splits(self):
+        self.assertEqual(list(morphlib.util.iter_trickle("foobarbazqux", 3)),
+                         [["f", "o", "o"], ["b", "a", "r"],
+                          ["b", "a", "z"], ["q", "u", "x"]])
+
+    def test_truncated_final_sequence(self):
+        self.assertEqual(list(morphlib.util.iter_trickle("barquux", 3)),
+                         [["b", "a", "r"], ["q", "u", "u"], ["x"]])