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 Group:

15 """ 16 Runs a subprocess in parallel, capturing it's output and providing non-blocking reads (well, at 17 least for the caller they appear non-blocking). 18 """

19 - def __init__(self):

20 self.output = Queue.Queue() 21 self.handles = [] 22 self.waiting = 0

23

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

25 """ 26 Adds a new process to this object. This process is run and the output collected. 27 28 @param cmd: the command to execute. This may be an array as passed to Popen, 29 or a string, which will be parsed by 'shlex.split' 30 @return: the handle to the process return from Popen 31 """ 32 cmd = _expand_cmd(cmd) 33 34 handle = subprocess.Popen( cmd, 35 shell = shell, 36 bufsize = 1, 37 stdout = subprocess.PIPE, 38 stderr = subprocess.STDOUT, 39 stdin = subprocess.PIPE, # needed to detach from calling terminal (other wacky things can happen) 40 ) 41 handle.group_output_done = False 42 self.handles.append( handle ) 43 44 # a thread is created to do blocking-read 45 self.waiting += 1 46 def block_read(): 47 try: 48 for line in iter( handle.stdout.readline, '' ): 49 self.output.put( ( handle, line ) ) 50 except: 51 pass 52 53 # To force return of any waiting read (and indicate this process is done 54 self.output.put( ( handle, None ) ) 55 handle.stdout.close() 56 self.waiting -= 1

57 58 block_thread = threading.Thread( target = block_read ) 59 block_thread.daemon = True 60 block_thread.start() 61 62 # kill child when parent dies 63 def premature_exit(): 64 try: 65 handle.terminate() 66 except: 67 pass # who cares why, we're exiting anyway (most likely since it is already terminated)

68 atexit.register( premature_exit ) 69 70 return handle 71

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

73 """ 74 Reads available lines from any of the running processes. If no lines are available now 75 it will wait until 'timeout' to read a line. If nothing is running the timeout is not waited 76 and the function simply returns. 77 78 When a process has been completed and all output has been read from it, a 79 variable 'group_ouput_done' will be set to True on the process handle. 80 81 @param timeout: how long to wait if there is nothing available now 82 @param max_lines: maximum number of lines to get at once 83 @return: An array of tuples of the form: 84 ( handle, line ) 85 There 'handle' was returned by 'run' and 'line' is the line which is read. 86 If no line is available an empty list is returned. 87 """ 88 lines = [] 89 try: 90 while len(lines) < max_lines: 91 handle, line = self.output.get_nowait() 92 # interrupt waiting if nothing more is expected 93 if line == None: 94 handle.group_output_done = True 95 if self.waiting == 0: 96 break 97 else: 98 lines.append( ( handle, line ) ) 99 return lines 100 101 except Queue.Empty: 102 # if nothing yet, then wait for something 103 if len(lines) > 0 or self.waiting == 0: 104 return lines 105 106 item = self.readline( timeout = timeout ) 107 if item != None: 108 lines.append( item ) 109 return lines

110

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

112 """ 113 Read a single line from any running process. 114 115 Note that this will end up blocking for timeout once all processes have completed. 116 'readlines' however can properly handle that situation and stop reading once 117 everything is complete. 118 119 @return: Tuple of ( handle, line ) or None if no output generated. 120 """ 121 try: 122 handle, line = self.output.get( timeout = timeout ) 123 if line == None: 124 handle.group_output_done = True 125 return None 126 return (handle, line) 127 except Queue.Empty: 128 return None

129

130 - def is_pending( self ):

131 """ 132 Determine if calling readlines would actually yield any output. This returns true 133 if there is a process running or there is data in the queue. 134 """ 135 if self.waiting > 0: 136 return True 137 return not self.output.empty()

138

139 - def count_running( self ):

140 """ 141 Return the number of processes still running. Note that although a process may 142 be finished there could still be output from it in the queue. You should use 'is_pending' 143 to determine if you should still be reading. 144 """ 145 count = 0 146 for handle in self.handles: 147 if handle.poll() == None: 148 count += 1 149 return count

150

151 - def get_exit_codes( self ):

152 """ 153 Return a list of all processes and their exit code. 154 155 @return: A list of tuples: 156 ( handle, exit_code ) 157 'handle' as returned from 'run' 158 'exit_code' of the process or None if it has not yet finished 159 """ 160 codes = [] 161 for handle in self.handles: 162 codes.append( ( handle, handle.poll() ) ) 163 return codes

164

165 - def clear_finished( self ):

166 """ 167 Remove all finished processes from the managed list. 168 """ 169 nhandles = [] 170 for handle in self.handles: 171 if not handle.group_output_done or handle.poll() == None: 172 nhandles.append( handle ) 173 self.handles = nhandles

174 175

176 -class BadExitCode(Exception):

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

178 Exception.__init__( self, 'subprocess-bad-exit-code: {}: {}'.format( exit_code, output[:1024] ) ) 179 self.exit_code = exit_code 180 self.output = output

181

182 -class Timeout(Exception):

183 - def __init__(self, output):

184 Exception.__init__( self, 'subprocess-timeout' ) 185 self.output = output

186

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

188 """ 189 Calls a subprocess and returns the output and optionally exit code. 190 191 @param encoding: convert output to unicode objects with this encoding, set to None to 192 get the raw output 193 @param check_exit_code: set to False to ignore the exit code, otherwise any non-zero 194 result will throw BadExitCode. 195 @param timeout: If specified only this amount of time (seconds) will be waited for 196 the subprocess to return 197 @return: If check_exit_code is False: list( output, exit_code ), else just the output 198 """ 199 cmd = _expand_cmd(cmd) 200 proc = subprocess.Popen( cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, 201 stdin = subprocess.PIPE, shell = shell ) 202 203 def decode(out): 204 if encoding != None: 205 return unicode( out, encoding ) 206 else: 207 return raw_out

208 209 if timeout == None: 210 raw_out, ignore_err = proc.communicate() 211 else: 212 # Read from subprocess in a thread so the main one can check for the timeout 213 outq = Queue.Queue() 214 def block_read(): 215 # collect as lines so if timeout we still have partial output 216 out = proc.stdout.read() 217 outq.put( out ) 218 219 block_thread = threading.Thread( target = block_read ) 220 block_thread.daemon = True 221 block_thread.start() 222 223 try: 224 raw_out = outq.get(True,timeout) 225 except Queue.Empty: 226 proc.terminate() 227 # wait again for partial output (process is terminated, so reading should end) 228 raw_out = outq.get() 229 raise Timeout( decode(raw_out) ) 230 231 out = decode(raw_out) 232 exit_code = proc.poll() 233 234 if check_exit_code: 235 if exit_code != 0: 236 raise BadExitCode( exit_code, out ) 237 return out 238 239 return ( out, proc.poll() ) 240

241 -def _expand_cmd(cmd):

242 if isinstance(cmd, basestring): 243 cmd = shlex.split(cmd) 244 return cmd

245

Source Code for Module shelljob.proc