# HG changeset patch
# User Jeff Hammel <k0scist@gmail.com>
# Date 1391383949 28800
# Node ID 8e23bbc9c1977fe7d823f859aa287385232bc51d
# Parent  5ce25399da67efcb6663151ca91de12253363a19
STUB: python/multiproc.py

diff -r 5ce25399da67 -r 8e23bbc9c197 python/multiproc.py
--- a/python/multiproc.py	Sun Feb 02 12:40:07 2014 -0800
+++ b/python/multiproc.py	Sun Feb 02 15:32:29 2014 -0800
@@ -18,6 +18,7 @@
 
     # http://docs.python.org/2/library/subprocess.html#popen-constructor
     defaults = {'bufsize': 1, # line buffered
+                'store_output': True, # store stdout
                 }
 
     def __init__(self, command, **kwargs):
@@ -33,16 +34,31 @@
             _kwargs['shell'] = isinstance(command, string)
 
         # output buffer
+        self.location = 0
         self.output_buffer = tempfile.SpooledTemporaryFile()
-        self.location = 0
-        self.output = ''
+        self.output = '' if _kwargs.pop('store_output') else None
         _kwargs['stdout'] = self.output_buffer
 
+        # runtime
+        self.start = time.time()
+        self.end = None
+
         # launch subprocess
-        self.start = time.time()
         subprocess.Popen.__init__(self, command, **_kwargs)
 
-    def wait(self, maxtime=None, sleep=1.):
+    def _finalize(self, process_output):
+        """internal function to finalize"""
+
+        # read final output
+        self.read(process_output)
+
+        # reset output buffer
+        self.output_buffer.seek(0)
+
+        # set end time
+        self.end = time.time()
+
+    def wait(self, maxtime=None, sleep=1., process_output=None):
         """
         maxtime -- timeout in seconds
         sleep -- number of seconds to sleep between polling
@@ -53,26 +69,31 @@
             curr_time = time.time()
             run_time = curr_time - self.start
             if run_time > maxtime:
-                # TODO: read from output
-                return process.kill()
+                self.kill()
+                self._finalize(process_output)
+                return
 
             # read from output buffer
-            read = self.read()
+            self.read(process_output)
 
             # naptime
             if sleep:
                 time.sleep(sleep)
 
-        # reset tempfile
-        output.seek(0)
+        # finalize
+        self._finalize()
 
         return self.returncode # set by ``.poll()``
 
-    def read(self):
+    def read(self, process_output=None):
         """read from the output buffer"""
+
         self.output_buffer.seek(self.location)
         read = self.output_buffer.read()
-        self.output += read
+        if self.output is not None:
+            self.output += read
+        if process_output:
+            process_output(read)
         self.location += len(read)
         return read
 
@@ -85,6 +106,13 @@
 
     __str__ = commandline
 
+    def runtime(self):
+        """returns time spent running or total runtime if completed"""
+
+        if self.end is None:
+            return self.end - self.start
+        return time.time() - self.start
+
 
 def main(args=sys.argv[1:]):
     """CLI"""
@@ -93,10 +121,8 @@
     progs = {'yes': ["yes"],
              'ping': ['ping', 'google.com']}
 
-
     # parse command line
-    usage = '%prog [options]'
-    parser = argparse.ArgumentParser(usage=usage, description=__doc__)
+    parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("-t", "--time", dest="time",
                         type=float, default=4.,
                         help="seconds to run for")
@@ -120,8 +146,13 @@
     # select program
     prog = progs[options.program]
 
+    # start process
     proc = Process(prog)
 
+    # callback for output processing
+    def process_output(output):
+        print output.upper()
+
     # # start the main subprocess loop
     # # TODO -> OO
     # output = tempfile.SpooledTemporaryFile()
@@ -143,7 +174,12 @@
     # # reset tempfile
     # output.seek(0)
 
-    n_lines = len(output.read().splitlines())
+    # wait for being done
+    proc.wait(maxtime=options.time, sleep=options.sleep, process_output=process_output)
+
+    # finalization
+    output = proc.output
+    n_lines = len(output.splitlines())
     print ("{}: {} lines".format(subprocess.list2cmdline(prog), n_lines))
 
 if __name__ == '__main__':