python/rox/processes.py

   1 """This module makes it easier to use other programs to process data.
   2
   3 The Process class provides the low-level interface, which you can extend
   4 by subclassing. Processes run in the background, so users can still work
   5 with your application while they are running. If you don't care about that,
   6 you might like to look at Python's builtin popen2 module.
   7
   8 The PipeThroughCommand class extends Process to provide an easy way to
   9 run other commands. It also, optionally, allows a stream of data to be fed
  10 in to the process's standard input, and can collect the output to another
  11 stream. Typical usage:
  12
  13 rox.processes.PipeThroughCommand(('echo', 'hello'), None, file('output', 'w')).wait()
  14
  15 This creates a new process, and execs 'echo hello' in it with output sent
  16 to the file 'output' (any file-like object can be used). The wait() runs a
  17 recursive mainloop, so that your application can still be used while the
  18 command runs, but the wait() itself doesn't return until the command
  19 completes.
  20
  21 Instead of using a tuple for the command, a string may be passed (eg, "echo
  22 hello"). In this case, the shell is used to interpret the command, allowing
  23 pipes, wildcards and so on. Be very careful of escaping in this case (think
  24 about filenames containing spaces, quotes, apostrophes, etc).
  25 """
  26
  27 from rox import g, saving
  28
  29 import os, sys, fcntl
  30 import signal
  31
  32 def _keep_on_exec(fd): fcntl.fcntl(fd, fcntl.F_SETFD, 0)
  33
  34 class ChildError(Exception):
  35         "Raised when the child process reports an error."
  36         def __init__(self, message):
  37                 Exception.__init__(self, message)
  38
  39 class ChildKilled(ChildError):
  40         "Raised when child died due to a call to the kill method."
  41         def __init__(self):
  42                 ChildError.__init__(self, "Operation aborted at user's request")
  43
  44 class Process:
  45         """This represents another process. You should subclass this
  46         and override the various methods. Use this when you want to
  47         run another process in the background, but still be able to
  48         communicate with it."""
  49         def __init__(self):
  50                 self.child = None
  51
  52         def start(self):
  53                 """Create the subprocess. Calls pre_fork() and forks.
  54                 The parent then calls parent_post_fork() and returns,
  55                 while the child calls child_post_fork() and then
  56                 child_run()."""
  57
  58                 assert self.child is None
  59
  60                 stderr_r = stderr_w = None
  61
  62                 try:
  63                         self.pre_fork()
  64                         stderr_r, stderr_w = os.pipe()
  65                         child = os.fork()
  66                 except:
  67                         if stderr_r: os.close(stderr_r)
  68                         if stderr_w: os.close(stderr_w)
  69                         self.start_error()
  70                         raise
  71
  72                 if child == 0:
  73                         # This is the child process
  74                         try:
  75                                 try:
  76                                         os.setpgid(0, 0)  # Start a new process group
  77                                         os.close(stderr_r)
  78
  79                                         if stderr_w != 2:
  80                                                 os.dup2(stderr_w, 2)
  81                                                 os.close(stderr_w)
  82
  83                                         self.child_post_fork()
  84                                         self.child_run()
  85                                         raise Exception('child_run() returned!')
  86                                 except:
  87                                         import traceback
  88                                         traceback.print_exc()
  89                         finally:
  90                                 os._exit(1)
  91                         assert 0
  92
  93                 self.child = child
  94
  95                 # This is the parent process
  96                 os.close(stderr_w)
  97                 self.err_from_child = stderr_r
  98
  99                 import gobject
 100                 if not hasattr(gobject, 'io_add_watch'):
 101                         self.tag = g.input_add_full(self.err_from_child,
 102                                         g.gdk.INPUT_READ, self._got_errors)
 103                 else:
 104                         self.tag = gobject.io_add_watch(self.err_from_child,
 105                                         gobject.IO_IN | gobject.IO_HUP | gobject.IO_ERR,
 106                                         self._got_errors)
 107
 108                 self.parent_post_fork()
 109
 110         def pre_fork(self):
 111                 """This is called in 'start' just before forking into
 112                 two processes. If you want to share a resource between
 113                 both processes (eg, a pipe), create it here.
 114                 Default method does nothing."""
 115
 116         def parent_post_fork(self):
 117                 """This is called in the parent after forking. Free the
 118                 child part of any resources allocated in pre_fork().
 119                 Also called if the fork or pre_fork() fails.
 120                 Default method does nothing."""
 121
 122         def child_post_fork(self):
 123                 """Called in the child after forking. Release the parent
 124                 part of any resources allocated in pre_fork().
 125                 Also called (in the parent) if the fork or pre_fork()
 126                 fails. Default method does nothing."""
 127
 128         def start_error(self):
 129                 """An error occurred before or during the fork (possibly
 130                 in pre_fork(). Clean up. Default method calls
 131                 parent_post_fork() and child_post_fork(). On returning,
 132                 the original exception will be raised."""
 133                 self.parent_post_fork()
 134                 self.child_post_fork()
 135
 136         def child_run(self):
 137                 """Called in the child process (after child_post_fork()).
 138                 Do whatever processing is required (perhaps exec another
 139                 process). If you don't exec, call os._exit(n) when done.
 140                 DO NOT make gtk calls in the child process, as it shares its
 141                 parent's connection to the X server until you exec()."""
 142                 os._exit(0)
 143
 144         def kill(self, sig = signal.SIGTERM):
 145                 """Send a signal to all processes in the child's process
 146                 group. The default, SIGTERM, requests all the processes
 147                 terminate. SIGKILL is more forceful."""
 148                 assert self.child is not None
 149                 os.kill(-self.child, sig)
 150
 151         def got_error_output(self, data):
 152                 """Read some characters from the child's stderr stream.
 153                 The default method copies to our stderr. Note that 'data'
 154                 isn't necessarily a complete line; it could be a single
 155                 character, or several lines, etc."""
 156                 sys.stderr.write(data)
 157
 158         def _got_errors(self, source, cond):
 159                 got = os.read(self.err_from_child, 100)
 160                 if got:
 161                         self.got_error_output(got)
 162                         return 1
 163
 164                 os.close(self.err_from_child)
 165                 g.input_remove(self.tag)
 166                 del self.tag
 167
 168                 pid, status = os.waitpid(self.child, 0)
 169                 self.child = None
 170                 self.child_died(status)
 171
 172         def child_died(self, status):
 173                 """Called when the child died (actually, when the child
 174                 closes its end of the stderr pipe). The child process has
 175                 already been reaped at this point; 'status' is the status
 176                 returned by os.waitpid."""
 177
 178 class PipeThroughCommand(Process):
 179         def __init__(self, command, src, dst):
 180                 """Execute 'command' with src as stdin and writing to stream
 181                 dst. If either stream is not a fileno() stream, temporary files
 182                 will be used as required.
 183                 Either stream may be None if input or output is not required.
 184                 Call the wait() method to wait for the command to finish.
 185                 'command' may be a string (passed to os.system) or a list (os.execvp).
 186                 """
 187
 188                 if src is not None and not hasattr(src, 'fileno'):
 189                         import shutil
 190                         new = _Tmp()
 191                         src.seek(0)
 192                         shutil.copyfileobj(src, new)
 193                         src = new
 194
 195                 Process.__init__(self)
 196
 197                 self.command = command
 198                 self.dst = dst
 199                 self.src = src
 200                 self.tmp_stream = None
 201
 202                 self.callback = None
 203                 self.killed = 0
 204                 self.errors = ""
 205
 206                 self.done = False       # bool or exception
 207                 self.waiting = False
 208
 209         def pre_fork(self):
 210                 # Output to 'dst' directly if it's a fileno stream. Otherwise,
 211                 # send output to a temporary file.
 212                 assert self.tmp_stream is None
 213
 214                 if self.dst:
 215                         if hasattr(self.dst, 'fileno'):
 216                                 self.dst.flush()
 217                                 self.tmp_stream = self.dst
 218                         else:
 219                                 self.tmp_stream = _Tmp()
 220
 221         def start_error(self):
 222                 self.tmp_stream = None
 223
 224         def child_run(self):
 225                 src = self.src
 226
 227                 if src:
 228                         os.dup2(src.fileno(), 0)
 229                         _keep_on_exec(0)
 230                         os.lseek(0, 0, 0)       # OpenBSD needs this, dunno why
 231                 if self.dst:
 232                         os.dup2(self.tmp_stream.fileno(), 1)
 233                         _keep_on_exec(1)
 234
 235                 # (basestr is python2.3 only)
 236                 if isinstance(self.command, str):
 237                         if os.system(self.command) == 0:
 238                                 os._exit(0)     # No error code or signal
 239                 else:
 240                         os.execvp(self.command[0], self.command)
 241                 os._exit(1)
 242
 243         def parent_post_fork(self):
 244                 if self.dst and self.tmp_stream is self.dst:
 245                         self.tmp_stream = None
 246
 247         def got_error_output(self, data):
 248                 self.errors += data
 249
 250         def check_errors(self, errors, status):
 251                 """Raise an exception here if errors (the string the child process wrote to stderr) or
 252                 status (the status from waitpid) seems to warrent it. It will be returned by wait()."""
 253                 if errors:
 254                         raise ChildError("Errors from command '%s':\n%s" % (str(self.command), errors))
 255                 raise ChildError("Command '%s' returned an error code (%d)!" % (str(self.command), status))
 256
 257         def child_died(self, status):
 258                 errors = self.errors.strip()
 259
 260                 if self.killed:
 261                         self.done = ChildKilled()
 262                 elif errors or status:
 263                         try:
 264                                 self.check_errors(errors, status)
 265                                 self.done = True
 266                         except Exception, e:
 267                                 self.done = e
 268                 else:
 269                         self.done = True
 270
 271                 assert self.done is True or isinstance(self.done, Exception)
 272
 273                 if self.done is True:
 274                         # Success
 275                         # If dst wasn't a fileno stream, copy from the temp file to it
 276                         if self.tmp_stream:
 277                                 self.tmp_stream.seek(0)
 278                                 self.dst.write(self.tmp_stream.read())
 279
 280                 self.tmp_stream = None
 281
 282                 if self.waiting:
 283                         assert self.done
 284                         self.waiting = False
 285                         g.mainquit()
 286
 287         def wait(self):
 288                 """Run a recursive mainloop until the command terminates.
 289                 Raises an exception on error."""
 290                 if self.child is None:
 291                         self.start()
 292                 self.waiting = True
 293                 while not self.done:
 294                         g.mainloop()
 295                 if self.done is not True:
 296                         raise self.done
 297
 298         def kill(self):
 299                 self.killed = 1
 300                 Process.kill(self)
 301
 302 def _Tmp(mode = 'w+b', suffix = '-tmp'):
 303         "Create a seekable, randomly named temp file (deleted automatically after use)."
 304         import tempfile
 305         try:
 306                 return tempfile.NamedTemporaryFile(mode, suffix = suffix)
 307         except:
 308                 # python2.2 doesn't have NamedTemporaryFile...
 309                 pass
 310
 311         import random
 312         name = tempfile.mktemp(`random.randint(1, 1000000)` + suffix)
 313
 314         fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
 315         tmp = tempfile.TemporaryFileWrapper(os.fdopen(fd, mode), name)
 316         tmp.name = name
 317         return tmp
 318
 319
 320 def _test():
 321         "Check that this module works."
 322
 323         def show():
 324                 error = sys.exc_info()[1]
 325                 print "(error reported was '%s')" % error
 326
 327         def pipe_through_command(command, src, dst): PipeThroughCommand(command, src, dst).wait()
 328
 329         print "Test _Tmp()..."
 330
 331         file = _Tmp()
 332         file.write('Hello')
 333         print >>file, ' ',
 334         file.flush()
 335         os.write(file.fileno(), 'World')
 336
 337         file.seek(0)
 338         assert file.read() == 'Hello World'
 339
 340         print "Test pipe_through_command():"
 341
 342         print "Try an invalid command..."
 343         try:
 344                 pipe_through_command('bad_command_1234', None, None)
 345                 assert 0
 346         except ChildError:
 347                 show()
 348         else:
 349                 assert 0
 350
 351         print "Try a valid command..."
 352         pipe_through_command('exit 0', None, None)
 353
 354         print "Writing to a non-fileno stream..."
 355         from cStringIO import StringIO
 356         a = StringIO()
 357         pipe_through_command('echo Hello', None, a)
 358         assert a.getvalue() == 'Hello\n'
 359
 360         print "Try with args..."
 361         a = StringIO()
 362         pipe_through_command(('echo', 'Hello'), None, a)
 363         assert a.getvalue() == 'Hello\n'
 364
 365         print "Reading from a stream to a StringIO..."
 366         file.seek(1)                    # (ignored)
 367         pipe_through_command('cat', file, a)
 368         assert a.getvalue() == 'Hello\nHello World'
 369
 370         print "Writing to a fileno stream..."
 371         file.seek(0)
 372         file.truncate(0)
 373         pipe_through_command('echo Foo', None, file)
 374         file.seek(0)
 375         assert file.read() == 'Foo\n'
 376
 377         print "Read and write fileno streams..."
 378         src = _Tmp()
 379         src.write('123')
 380         src.seek(0)
 381         file.seek(0)
 382         file.truncate(0)
 383         pipe_through_command('cat', src, file)
 384         file.seek(0)
 385         assert file.read() == '123'
 386
 387         print "Detect non-zero exit value..."
 388         try:
 389                 pipe_through_command('exit 1', None, None)
 390         except ChildError:
 391                 show()
 392         else:
 393                 assert 0
 394
 395         print "Detect writes to stderr..."
 396         try:
 397                 pipe_through_command('echo one >&2; sleep 2; echo two >&2', None, None)
 398         except ChildError:
 399                 show()
 400         else:
 401                 assert 0
 402
 403         print "Check tmp file is deleted..."
 404         name = file.name
 405         assert os.path.exists(name)
 406         file = None
 407         assert not os.path.exists(name)
 408
 409         print "Check we can kill a runaway proces..."
 410         ptc = PipeThroughCommand('sleep 100; exit 1', None, None)
 411         def stop():
 412                 ptc.kill()
 413         g.timeout_add(2000, stop)
 414         try:
 415                 ptc.wait()
 416                 assert 0
 417         except ChildKilled:
 418                 pass
 419
 420         print "All tests passed!"
 421
 422 if __name__ == '__main__':
 423         _test()