code.delx.au - bg-scripts/blob - lib/commands_async.py

   1 #! python
   2
   3 """
   4 A small utility that provides similar functionality to the commands module, but
   5 allows you to get the output from multiple processes at the same time
   6 """
   7
   8 __author__ = "Greg Darke"
   9
  10 import subprocess, fcntl, os
  11 from select import select
  12 try:
  13         import cStringIO as _StringIO
  14 except ImportError:
  15         import StringIO as _StringIO
  16
  17 class CommandRunner(object):
  18         def __init__(self):
  19                 self._outputs = {}
  20                 self._processes = {}
  21                 self._fds = []
  22
  23         def _executeCommand(self, cmd):
  24                 """Execute the command"""
  25                 output = _StringIO.StringIO()
  26                 isShell = isinstance(cmd, str)
  27                 process = subprocess.Popen(args = cmd, shell = isShell,
  28                                 stdout = subprocess.PIPE, stderr = subprocess.STDOUT)
  29
  30                 # Turn blocking off
  31                 flags = fcntl.fcntl(process.stdout, fcntl.F_GETFL) | os.O_NONBLOCK
  32                 fcntl.fcntl(process.stdout, fcntl.F_SETFL, flags)
  33
  34                 return (output, process)
  35
  36         def executeCommand(self, cmd):
  37                 """Executes a command, but does not return anything"""
  38                 output, process = self._executeCommand(cmd)
  39                 self._outputs[process.stdout] = output
  40                 self._processes[cmd] = process
  41                 self._fds.append(process.stdout)
  42
  43         def _waitLoop(self):
  44                 count = 0
  45                 while self._fds:
  46                         # While there are still some processes left
  47                         inputReady, outputReady, exceptReady = \
  48                                         select(self._fds, [], self._fds)
  49
  50                         for fd in inputReady:
  51                                 data = fd.read()
  52                                 if not data:
  53                                         self._fds.remove(fd)
  54                                         continue
  55                                 self._outputs[fd].write(data)
  56
  57                         for fd in exceptReady:
  58                                 self._fds.remove(fd)
  59
  60         def waitForCompletion(self):
  61                 """Waits for all of the running processes to finish"""
  62                 self._waitLoop()
  63
  64         def getOutputs(self):
  65                 """Returns a dictionay containing the command as the key, and a string (of the output) as the value"""
  66                 outputs = dict((cmd, self._outputs[process.stdout].getvalue()) for cmd, process in self._processes.items())
  67                 return outputs