Allow runcmd callbacks to mangle the collected data

2 files changed, 18 insertions, 8 deletions

diff --git a/cliapp/runcmd.py b/cliapp/runcmd.py
index dc33011..ee9698e 100644
--- a/cliapp/runcmd.py
+++ b/cliapp/runcmd.py
@@ -202,8 +202,10 @@ def _run_pipeline(procs, feed_stdin, pipe_stdin, pipe_stdout, pipe_stderr,
             data = procs[-1].stdout.read(io_size)
             if data:
                 logging.debug('calling stdout callback: %r', stdout_callback)
-                stdout_callback(data)
-                out.append(data)
+                data_new = stdout_callback(data)
+                if data_new is None:
+                    data_new = data
+                out.append(data_new)
             else:
                 stdout_eof = True
 
@@ -211,8 +213,10 @@ def _run_pipeline(procs, feed_stdin, pipe_stdin, pipe_stdout, pipe_stderr,
             data = procs[-1].stderr.read(io_size)
             if data:
                 logging.debug('calling stderr callback: %r', stderr_callback)
-                stderr_callback(data)
-                err.append(data)
+                data_new = stderr_callback(data)
+                if data_new is None:
+                    data_new = data
+                err.append(data_new)
             else:
                 stderr_eof = True
 
diff --git a/cliapp/runcmd_tests.py b/cliapp/runcmd_tests.py
index 0e8a090..3346830 100644
--- a/cliapp/runcmd_tests.py
+++ b/cliapp/runcmd_tests.py
@@ -130,13 +130,16 @@ class RuncmdTests(unittest.TestCase):
         def logger(s):
             msgs.append(s)
 
+            # We return a string to allow the callback to mangle
+            # the data being returned.
+            return 'foo'
+
         test_input = 'hello fox'
         exit, out, err = cliapp.runcmd_unchecked(['echo', '-n', test_input],
                                                  stdout_callback=logger)
 
-        self.assertEqual(test_input, out)
-        self.assertEqual(len(msgs), 1)
-        self.assertEqual(out, msgs[0])
+        self.assertEqual(out, 'foo')
+        self.assertEqual(msgs, [test_input])
 
     def test_runcmd_calls_stderr_callback_when_msg_on_stderr(self):
         msgs = []
@@ -144,6 +147,10 @@ class RuncmdTests(unittest.TestCase):
         def logger(s):
             msgs.append(s)
 
+            # We return None to signal that the data should not be
+            # mangled.
+            return None
+
         exit, out, err = cliapp.runcmd_unchecked(['ls', 'nosuchthing'],
                                                  stderr_callback=logger)
 
@@ -168,4 +175,3 @@ class ShellQuoteTests(unittest.TestCase):
 
     def test_quotes_single_quote(self):
         self.assertEqual(cliapp.shell_quote("'"), '"\'"')
-