shelljob.proc

1 # Subprocess containers 2 """ 3 A mechanism to run subprocesses asynchronously and with non-blocking read. 4 """ 5 import atexit 6 import multiprocessing 7 import Queue 8 import shlex 9 import subprocess 10 import sys 11 import threading 12 import time 13

14 -class WrapException(Exception):

15 - def __init__(self, exc, msg):

16 """ 17 Assumes being created within an exception handler. 18 Try to adopt Python 3 names for the bits 19 """ 20 self.__traceback__ = sys.exc_info()[2] 21 self.__cause__ = exc 22 self.message = msg 23 super(WrapException,self).__init__( msg, exc )

24

25 - def do_raise(self):

26 raise self, None, self.__traceback__

27

28 -class CommandException(WrapException):

29 - def __init__(self, exc, msg ):

30 super(CommandException,self).__init__( exc, msg )

31

32 -class Group:

33 """ 34 Runs a subprocess in parallel, capturing it's output and providing non-blocking reads (well, at 35 least for the caller they appear non-blocking). 36 """

37 - def __init__(self):

38 self.output = Queue.Queue() 39 self.handles = [] 40 self.waiting = 0

41

42 - def run( self, cmd, shell = False ):

43 """ 44 Adds a new process to this object. This process is run and the output collected. 45 46 @param cmd: the command to execute. This may be an array as passed to Popen, 47 or a string, which will be parsed by 'shlex.split' 48 @return: the handle to the process return from Popen 49 """ 50 try: 51 return self._run_impl( cmd, shell ) 52 except Exception, e: 53 CommandException( e, "Group.run '{}' failed".format( cmd ) ).do_raise()

54

55 - def _run_impl( self, cmd, shell ):

56 cmd = _expand_cmd(cmd) 57 58 handle = subprocess.Popen( cmd, 59 shell = shell, 60 bufsize = 1, 61 stdout = subprocess.PIPE, 62 stderr = subprocess.STDOUT, 63 stdin = subprocess.PIPE, # needed to detach from calling terminal (other wacky things can happen) 64 ) 65 handle.group_output_done = False 66 self.handles.append( handle ) 67 68 # a thread is created to do blocking-read 69 self.waiting += 1 70 def block_read(): 71 try: 72 for line in iter( handle.stdout.readline, '' ): 73 self.output.put( ( handle, line ) ) 74 except: 75 pass 76 77 # To force return of any waiting read (and indicate this process is done 78 self.output.put( ( handle, None ) ) 79 handle.stdout.close() 80 self.waiting -= 1

81 82 block_thread = threading.Thread( target = block_read ) 83 block_thread.daemon = True 84 block_thread.start() 85 86 # kill child when parent dies 87 def premature_exit(): 88 try: 89 handle.terminate() 90 except: 91 pass # who cares why, we're exiting anyway (most likely since it is already terminated)

92 atexit.register( premature_exit ) 93 94 return handle 95

96 - def readlines( self, max_lines = 1000, timeout = 2.0 ):

97 """ 98 Reads available lines from any of the running processes. If no lines are available now 99 it will wait until 'timeout' to read a line. If nothing is running the timeout is not waited 100 and the function simply returns. 101 102 When a process has been completed and all output has been read from it, a 103 variable 'group_ouput_done' will be set to True on the process handle. 104 105 @param timeout: how long to wait if there is nothing available now 106 @param max_lines: maximum number of lines to get at once 107 @return: An array of tuples of the form: 108 ( handle, line ) 109 There 'handle' was returned by 'run' and 'line' is the line which is read. 110 If no line is available an empty list is returned. 111 """ 112 lines = [] 113 try: 114 while len(lines) < max_lines: 115 handle, line = self.output.get_nowait() 116 # interrupt waiting if nothing more is expected 117 if line == None: 118 handle.group_output_done = True 119 if self.waiting == 0: 120 break 121 else: 122 lines.append( ( handle, line ) ) 123 return lines 124 125 except Queue.Empty: 126 # if nothing yet, then wait for something 127 if len(lines) > 0 or self.waiting == 0: 128 return lines 129 130 item = self.readline( timeout = timeout ) 131 if item != None: 132 lines.append( item ) 133 return lines

134

135 - def readline( self, timeout = 2.0 ):

136 """ 137 Read a single line from any running process. 138 139 Note that this will end up blocking for timeout once all processes have completed. 140 'readlines' however can properly handle that situation and stop reading once 141 everything is complete. 142 143 @return: Tuple of ( handle, line ) or None if no output generated. 144 """ 145 try: 146 handle, line = self.output.get( timeout = timeout ) 147 if line == None: 148 handle.group_output_done = True 149 return None 150 return (handle, line) 151 except Queue.Empty: 152 return None

153

154 - def is_pending( self ):

155 """ 156 Determine if calling readlines would actually yield any output. This returns true 157 if there is a process running or there is data in the queue. 158 """ 159 if self.waiting > 0: 160 return True 161 return not self.output.empty()

162

163 - def count_running( self ):

164 """ 165 Return the number of processes still running. Note that although a process may 166 be finished there could still be output from it in the queue. You should use 'is_pending' 167 to determine if you should still be reading. 168 """ 169 count = 0 170 for handle in self.handles: 171 if handle.poll() == None: 172 count += 1 173 return count

174

175 - def get_exit_codes( self ):

176 """ 177 Return a list of all processes and their exit code. 178 179 @return: A list of tuples: 180 ( handle, exit_code ) 181 'handle' as returned from 'run' 182 'exit_code' of the process or None if it has not yet finished 183 """ 184 codes = [] 185 for handle in self.handles: 186 codes.append( ( handle, handle.poll() ) ) 187 return codes

188

189 - def clear_finished( self ):

190 """ 191 Remove all finished processes from the managed list. 192 """ 193 nhandles = [] 194 for handle in self.handles: 195 if not handle.group_output_done or handle.poll() == None: 196 nhandles.append( handle ) 197 self.handles = nhandles

198 199

200 -class BadExitCode(Exception):

201 - def __init__(self, exit_code, output):

202 Exception.__init__( self, 'subprocess-bad-exit-code: {}: {}'.format( exit_code, output[:1024] ) ) 203 self.exit_code = exit_code 204 self.output = output

205

206 -class Timeout(Exception):

207 - def __init__(self, output):

208 Exception.__init__( self, 'subprocess-timeout' ) 209 self.output = output

210

211 -def call( cmd, encoding = 'utf-8', shell = False, check_exit_code = True, timeout = None ):

212 """ 213 Calls a subprocess and returns the output and optionally exit code. 214 215 @param encoding: convert output to unicode objects with this encoding, set to None to 216 get the raw output 217 @param check_exit_code: set to False to ignore the exit code, otherwise any non-zero 218 result will throw BadExitCode. 219 @param timeout: If specified only this amount of time (seconds) will be waited for 220 the subprocess to return 221 @return: If check_exit_code is False: list( output, exit_code ), else just the output 222 """ 223 cmd = _expand_cmd(cmd) 224 proc = subprocess.Popen( cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, 225 stdin = subprocess.PIPE, shell = shell ) 226 227 def decode(out): 228 if encoding != None: 229 return unicode( out, encoding ) 230 else: 231 return raw_out

232 233 if timeout == None: 234 raw_out, ignore_err = proc.communicate() 235 else: 236 # Read from subprocess in a thread so the main one can check for the timeout 237 outq = Queue.Queue() 238 def block_read(): 239 out = proc.stdout.read() 240 # wait before pushing, occassionally read returns prior to process terminating, 241 # thus "poll" would return None 242 proc.wait() 243 outq.put( out ) 244 245 block_thread = threading.Thread( target = block_read ) 246 block_thread.daemon = True 247 block_thread.start() 248 249 try: 250 raw_out = outq.get(True,timeout) 251 except Queue.Empty: 252 proc.terminate() 253 # wait again for partial output (process is terminated, so reading should end) 254 raw_out = outq.get() 255 raise Timeout( decode(raw_out) ) 256 257 out = decode(raw_out) 258 exit_code = proc.poll() 259 260 if check_exit_code: 261 if exit_code != 0: 262 raise BadExitCode( exit_code, out ) 263 return out 264 265 return ( out, proc.poll() ) 266

267 -def _expand_cmd(cmd):

268 if isinstance(cmd, basestring): 269 cmd = shlex.split(cmd) 270 return cmd

271

Source Code for Module shelljob.proc