support.py

   1 #!/usr/bin/env python
   2
   3 import findrox
   4 from rox import g, saving
   5 import rox
   6 import fcntl
   7
   8 try:
   9         from rox import processes
  10 except ImportError:
  11         rox.croak('Sorry, this version of Archive requires ROX-Lib 1.9.3 or later')
  12
  13 import sys, os
  14
  15 class ChildError(Exception):
  16         "Raised when the child process reports an error."
  17
  18 class ChildKilled(saving.AbortSave):
  19         "Raised when child died due to calling the kill method."
  20         def __init__(self):
  21                 saving.AbortSave.__init__(self, "Operation aborted at user's request")
  22
  23 def escape(text):
  24         """Return text with \ and ' escaped"""
  25         return text.replace("\\", "\\\\").replace("'", "\\'")
  26
  27 def Tmp(mode = 'w+b'):
  28         "Create a seekable, randomly named temp file (deleted automatically after use)."
  29         import tempfile
  30         try:
  31                 return tempfile.NamedTemporaryFile(mode, suffix = '-archive')
  32         except:
  33                 # python2.2 doesn't have NamedTemporaryFile...
  34                 pass
  35
  36         import random
  37         name = tempfile.mktemp(`random.randint(1, 1000000)` + '-archive')
  38
  39         fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
  40         tmp = tempfile.TemporaryFileWrapper(os.fdopen(fd, mode), name)
  41         tmp.name = name
  42         return tmp
  43
  44 def keep_on_exec(fd):
  45         fcntl.fcntl(fd, fcntl.F_SETFD, 0)
  46
  47 class PipeThroughCommand(processes.Process):
  48         def __init__(self, command, src, dst):
  49                 """Execute 'command' with src as stdin and writing to stream
  50                 dst. src must be a fileno() stream, but dst need not be.
  51                 Either stream may be None if input or output is not required.
  52                 Call the wait() method to wait for the command to finish."""
  53
  54                 assert src is None or hasattr(src, 'fileno')
  55
  56                 processes.Process.__init__(self)
  57
  58                 self.command = command
  59                 self.dst = dst
  60                 self.src = src
  61                 self.tmp_stream = None
  62
  63                 self.callback = None
  64                 self.killed = 0
  65                 self.errors = ""
  66
  67                 self.start()
  68
  69         def pre_fork(self):
  70                 # Output to 'dst' directly if it's a fileno stream. Otherwise,
  71                 # send output to a temporary file.
  72                 assert self.tmp_stream is None
  73
  74                 if self.dst:
  75                         if hasattr(self.dst, 'fileno'):
  76                                 self.dst.flush()
  77                                 self.tmp_stream = self.dst
  78                         else:
  79                                 self.tmp_stream = Tmp()
  80
  81         def start_error(self):
  82                 """Clean up effects of pre_fork()."""
  83                 self.tmp_stream = None
  84
  85         def child_run(self):
  86                 src = self.src
  87
  88                 if src:
  89                         os.dup2(src.fileno(), 0)
  90                         keep_on_exec(0)
  91                         os.lseek(0, 0, 0)       # OpenBSD needs this, dunno why
  92                 if self.dst:
  93                         os.dup2(self.tmp_stream.fileno(), 1)
  94                         keep_on_exec(1)
  95
  96                 if os.system(self.command) == 0:
  97                         os._exit(0)     # No error code or signal
  98                 os._exit(1)
  99
 100         def parent_post_fork(self):
 101                 if self.dst and self.tmp_stream is self.dst:
 102                         self.tmp_stream = None
 103
 104         def got_error_output(self, data):
 105                 self.errors += data
 106
 107         def child_died(self, status):
 108                 errors = self.errors.strip()
 109
 110                 err = None
 111
 112                 if self.killed:
 113                         err = ChildKilled
 114                 elif errors:
 115                         err = ChildError("Errors from command '%s':\n%s" % (self.command, errors))
 116                 elif status != 0:
 117                         err = ChildError("Command '%s' returned an error code!" % self.command)
 118
 119                 # If dst wasn't a fileno stream, copy from the temp file to it
 120                 if not err and self.tmp_stream:
 121                         self.tmp_stream.seek(0)
 122                         self.dst.write(self.tmp_stream.read())
 123                 self.tmp_stream = None
 124
 125                 self.callback(err)
 126
 127         def wait(self):
 128                 """Run a recursive mainloop until the command terminates.
 129                 Raises an exception on error."""
 130                 done = []
 131                 def set_done(exception):
 132                         done.append(exception)
 133                         g.mainquit()
 134                 self.callback = set_done
 135                 while not done:
 136                         g.mainloop()
 137                 exception, = done
 138                 if exception:
 139                         raise exception
 140
 141         def kill(self):
 142                 self.killed = 1
 143                 processes.Process.kill(self)
 144
 145 def test():
 146         "Check that this module works."
 147
 148         def show():
 149                 error = sys.exc_info()[1]
 150                 print "(error reported was '%s')" % error
 151
 152         def pipe_through_command(command, src, dst): PipeThroughCommand(command, src, dst).wait()
 153
 154         print "Test escape()..."
 155
 156         assert escape('''  a test ''') == '  a test '
 157         assert escape('''  "a's test" ''') == '''  "a\\'s test" '''
 158         assert escape('''  "a\\'s test" ''') == '''  "a\\\\\\'s test" '''
 159
 160         print "Test Tmp()..."
 161
 162         file = Tmp()
 163         file.write('Hello')
 164         print >>file, ' ',
 165         file.flush()
 166         os.write(file.fileno(), 'World')
 167
 168         file.seek(0)
 169         assert file.read() == 'Hello World'
 170
 171         print "Test pipe_through_command():"
 172
 173         print "Try an invalid command..."
 174         try:
 175                 pipe_through_command('bad_command_1234', None, None)
 176                 assert 0
 177         except ChildError:
 178                 show()
 179         else:
 180                 assert 0
 181
 182         print "Try a valid command..."
 183         pipe_through_command('exit 0', None, None)
 184
 185         print "Writing to a non-fileno stream..."
 186         from cStringIO import StringIO
 187         a = StringIO()
 188         pipe_through_command('echo Hello', None, a)
 189         assert a.getvalue() == 'Hello\n'
 190
 191         print "Reading from a stream to a StringIO..."
 192         file.seek(1)
 193         pipe_through_command('cat', file, a)
 194         assert a.getvalue() == 'Hello\nello World'
 195
 196         print "Writing to a fileno stream..."
 197         file.seek(0)
 198         file.truncate(0)
 199         pipe_through_command('echo Foo', None, file)
 200         file.seek(0)
 201         assert file.read() == 'Foo\n'
 202
 203         print "Read and write fileno streams..."
 204         src = Tmp()
 205         src.write('123')
 206         src.seek(0)
 207         file.seek(0)
 208         file.truncate(0)
 209         pipe_through_command('cat', src, file)
 210         file.seek(0)
 211         assert file.read() == '123'
 212
 213         print "Detect non-zero exit value..."
 214         try:
 215                 pipe_through_command('exit 1', None, None)
 216         except ChildError:
 217                 show()
 218         else:
 219                 assert 0
 220
 221         print "Detect writes to stderr..."
 222         try:
 223                 pipe_through_command('echo one >&2; sleep 2; echo two >&2', None, None)
 224         except ChildError:
 225                 show()
 226         else:
 227                 assert 0
 228
 229         print "Check tmp file is deleted..."
 230         name = file.name
 231         assert os.path.exists(name)
 232         file = None
 233         assert not os.path.exists(name)
 234
 235         print "Check we can kill a runaway proces..."
 236         ptc = PipeThroughCommand('sleep 100; exit 1', None, None)
 237         def stop():
 238                 ptc.kill()
 239         g.timeout_add(2000, stop)
 240         try:
 241                 ptc.wait()
 242                 assert 0
 243         except ChildKilled:
 244                 pass
 245
 246         print "All tests passed!"
 247
 248 if __name__ == '__main__':
 249         test()