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         def show():
  94                 error = sys.exc_info()[1]
  95                 print "(error reported was '%s')" % error
  96
  97         print "Test escape()..."
  98
  99         assert escape('''  a test ''') == '  a test '
 100         assert escape('''  "a's test" ''') == '''  "a\\'s test" '''
 101         assert escape('''  "a\\'s test" ''') == '''  "a\\\\\\'s test" '''
 102
 103         print "Test Tmp()..."
 104
 105         file = Tmp()
 106         file.write('Hello')
 107         print >>file, ' ',
 108         file.flush()
 109         os.write(file.fileno(), 'World')
 110
 111         file.seek(0)
 112         assert file.read() == 'Hello World'
 113
 114         print "Test pipe_through_command():"
 115
 116         print "Try an invalid command..."
 117         try:
 118                 pipe_through_command('bad_command_1234', None, None)
 119                 assert 0
 120         except ChildError:
 121                 show()
 122         else:
 123                 assert 0
 124
 125         print "Try a valid command..."
 126         pipe_through_command('exit 0', None, None)
 127
 128         print "Writing to a non-fileno stream..."
 129         from cStringIO import StringIO
 130         a = StringIO()
 131         pipe_through_command('echo Hello', None, a)
 132         assert a.getvalue() == 'Hello\n'
 133
 134         print "Reading from a stream to a StringIO..."
 135         file.seek(1)
 136         pipe_through_command('cat', file, a)
 137         assert a.getvalue() == 'Hello\nello World'
 138
 139         print "Writing to a fileno stream..."
 140         file.seek(0)
 141         file.truncate(0)
 142         pipe_through_command('echo Foo', None, file)
 143         file.seek(0)
 144         assert file.read() == 'Foo\n'
 145
 146         print "Read and write fileno streams..."
 147         src = Tmp()
 148         src.write('123')
 149         src.seek(0)
 150         file.seek(0)
 151         file.truncate(0)
 152         pipe_through_command('cat', src, file)
 153         file.seek(0)
 154         assert file.read() == '123'
 155
 156         print "Detect non-zero exit value..."
 157         try:
 158                 pipe_through_command('exit 1', None, None)
 159         except ChildError:
 160                 show()
 161         else:
 162                 assert 0
 163
 164         print "Detect writes to stderr..."
 165         try:
 166                 pipe_through_command('echo one >&2; sleep 2; echo two >&2', None, None)
 167         except ChildError:
 168                 show()
 169         else:
 170                 assert 0
 171
 172         print "All tests passed!"
 173
 174 if __name__ == '__main__':
 175         test()