support.py

   1 #!/usr/bin/env python
   2
   3 import sys, os
   4
   5 class ChildError(Exception):
   6         "Raised when the child process reports an error."
   7         pass
   8
   9 def escape(text):
  10         """Return text with \ and ' escaped"""
  11         return text.replace("\\", "\\\\").replace("'", "\\'")
  12
  13 def Tmp():
  14         "Create a Temporary file with a random name."
  15         import random
  16         name = `random.randint(1, 1000000)` + '-archive'
  17         import tempfile
  18         return tempfile.TemporaryFile(suffix = name)
  19
  20 def pipe_through_command(command, src, dst):
  21         """Execute 'command' with src as stdin (if any) and writing to
  22         stream dst. src must be a fileno() stream, but dst need not be.
  23         Either stream may be None if input or output is not required.
  24         Raises an Exception with a suitable message if the command
  25         writes to stderr or returns a non-zero exit status."""
  26
  27         assert src is None or hasattr(src, 'fileno')
  28
  29         # Output to 'dst' directly if it's a fileno stream. Otherwise,
  30         # send output to a temporary file.
  31         if dst:
  32                 if hasattr(dst, 'fileno'):
  33                         dst.flush()
  34                         tmp_stream = dst
  35                 else:
  36                         tmp_stream = Tmp()
  37                 fd = tmp_stream.fileno()
  38         else:
  39                 fd = -1
  40
  41         # Create a pipe to collect stderr from child
  42         stderr_r, stderr_w = os.pipe()
  43         try:
  44                 child = os.fork()
  45         except:
  46                 os.close(stderr_r)
  47                 os.close(stderr_w)
  48                 raise
  49
  50         if child == 0:
  51                 # This is the child process
  52                 try:
  53                         os.close(stderr_r)
  54                         os.dup2(stderr_w, 2)
  55                         if src:
  56                                 os.dup2(src.fileno(), 0)
  57                         if fd != -1:
  58                                 os.dup2(fd, 1)
  59                         if os.system(command) == 0:
  60                                 os._exit(0)     # No error code or signal
  61                 finally:
  62                         os._exit(1)
  63                 assert 0
  64
  65         # This is the parent process
  66         os.close(stderr_w)
  67
  68         # Read until child closes stderr (by quitting)
  69         errors = ""
  70         while 1:
  71                 got = os.read(stderr_r, 100)
  72                 if not got:
  73                         break
  74                 errors += got
  75         errors = errors.strip()
  76
  77         # Reap zombie
  78         pid, status = os.waitpid(child, 0)
  79
  80         if errors:
  81                 raise ChildError('Errors from command %s:\n%s' % (command, errors))
  82         if status != 0:
  83                 raise ChildError("Command '%s' returned an error code!" % command)
  84
  85         # If dst wasn't a fileno stream, copy from the temp file to it
  86         if dst and tmp_stream is not dst:
  87                 tmp_stream.seek(0)
  88                 dst.write(tmp_stream.read())
  89
  90 def test():
  91         "Check that this module works."
  92
  93         print "Test escape()..."
  94
  95         assert escape('''  a test ''') == '  a test '
  96         assert escape('''  "a's test" ''') == '''  "a\\'s test" '''
  97         assert escape('''  "a\\'s test" ''') == '''  "a\\\\\\'s test" '''
  98
  99         print "Test Tmp()..."
 100
 101         file = Tmp()
 102         file.write('Hello')
 103         print >>file, ' ',
 104         file.flush()
 105         os.write(file.fileno(), 'World')
 106
 107         file.seek(0)
 108         assert file.read() == 'Hello World'
 109
 110         print "Test pipe_through_command()..."
 111
 112         print "  invalid command..."
 113         try:
 114                 pipe_through_command('bad_command_1234', None, None)
 115                 assert 0
 116         except ChildError:
 117                 pass
 118         else:
 119                 assert 0
 120
 121         print "Writing to a non-fileno stream..."
 122         from cStringIO import StringIO
 123         a = StringIO()
 124         pipe_through_command('echo Hello', None, a)
 125         assert a.getvalue() == 'Hello\n'
 126
 127         print "Reading from a stream..."
 128         file.seek(1)
 129         pipe_through_command('cat', file, a)
 130         assert a.getvalue() == 'Hello\nello World'
 131
 132         print "Writing to a fileno stream..."
 133         file.seek(0)
 134         file.truncate(0)
 135         pipe_through_command('echo Foo', None, file)
 136         file.seek(0)
 137         assert file.read() == 'Foo\n'
 138
 139         print "Detect non-zero exit value..."
 140         try:
 141                 pipe_through_command('exit 1', None, None)
 142         except ChildError:
 143                 pass
 144         else:
 145                 assert 0
 146
 147         print "Detect writes to stderr..."
 148         try:
 149                 pipe_through_command('echo one >&2; sleep 2; echo two >&2', None, None)
 150         except ChildError:
 151                 pass
 152         else:
 153                 assert 0
 154
 155         print "All tests passed!"
 156
 157 if __name__ == '__main__':
 158         test()