support.py

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