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