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(mode = 'w+b'):
  14         "Create a seekable, randomly named temp file (deleted automatically after use)."
  15         import tempfile
  16         import random
  17         name = tempfile.mktemp(`random.randint(1, 1000000)` + '-archive')
  18
  19         fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
  20         tmp = tempfile.TemporaryFileWrapper(os.fdopen(fd, mode), name)
  21         tmp.name = name
  22         return tmp
  23
  24 def pipe_through_command(command, src, dst):
  25         """Execute 'command' with src as stdin (if any) and writing to
  26         stream dst. src must be a fileno() stream, but dst need not be.
  27         Either stream may be None if input or output is not required.
  28         Raises an Exception with a suitable message if the command
  29         writes to stderr or returns a non-zero exit status."""
  30
  31         assert src is None or hasattr(src, 'fileno')
  32
  33         # Output to 'dst' directly if it's a fileno stream. Otherwise,
  34         # send output to a temporary file.
  35         if dst:
  36                 if hasattr(dst, 'fileno'):
  37                         dst.flush()
  38                         tmp_stream = dst
  39                 else:
  40                         tmp_stream = Tmp()
  41                 fd = tmp_stream.fileno()
  42         else:
  43                 fd = -1
  44
  45         # Create a pipe to collect stderr from child
  46         stderr_r, stderr_w = os.pipe()
  47         try:
  48                 child = os.fork()
  49         except:
  50                 os.close(stderr_r)
  51                 os.close(stderr_w)
  52                 raise
  53
  54         if child == 0:
  55                 # This is the child process
  56                 try:
  57                         os.close(stderr_r)
  58                         os.dup2(stderr_w, 2)
  59                         if src:
  60                                 os.dup2(src.fileno(), 0)
  61                         if fd != -1:
  62                                 os.dup2(fd, 1)
  63                         if os.system(command) == 0:
  64                                 os._exit(0)     # No error code or signal
  65                 finally:
  66                         os._exit(1)
  67                 assert 0
  68
  69         # This is the parent process
  70         os.close(stderr_w)
  71
  72         # Read until child closes stderr (by quitting)
  73         errors = ""
  74         while 1:
  75                 got = os.read(stderr_r, 100)
  76                 if not got:
  77                         break
  78                 errors += got
  79         errors = errors.strip()
  80
  81         # Reap zombie
  82         pid, status = os.waitpid(child, 0)
  83
  84         if errors:
  85                 raise ChildError("Errors from command '%s':\n%s" % (command, errors))
  86         if status != 0:
  87                 raise ChildError("Command '%s' returned an error code!" % command)
  88
  89         # If dst wasn't a fileno stream, copy from the temp file to it
  90         if dst and tmp_stream is not dst:
  91                 tmp_stream.seek(0)
  92                 dst.write(tmp_stream.read())
  93
  94 def test():
  95         "Check that this module works."
  96
  97         def show():
  98                 error = sys.exc_info()[1]
  99                 print "(error reported was '%s')" % error
 100
 101         print "Test escape()..."
 102
 103         assert escape('''  a test ''') == '  a test '
 104         assert escape('''  "a's test" ''') == '''  "a\\'s test" '''
 105         assert escape('''  "a\\'s test" ''') == '''  "a\\\\\\'s test" '''
 106
 107         print "Test Tmp()..."
 108
 109         file = Tmp()
 110         file.write('Hello')
 111         print >>file, ' ',
 112         file.flush()
 113         os.write(file.fileno(), 'World')
 114
 115         file.seek(0)
 116         assert file.read() == 'Hello World'
 117
 118         print "Test pipe_through_command():"
 119
 120         print "Try an invalid command..."
 121         try:
 122                 pipe_through_command('bad_command_1234', None, None)
 123                 assert 0
 124         except ChildError:
 125                 show()
 126         else:
 127                 assert 0
 128
 129         print "Try a valid command..."
 130         pipe_through_command('exit 0', None, None)
 131
 132         print "Writing to a non-fileno stream..."
 133         from cStringIO import StringIO
 134         a = StringIO()
 135         pipe_through_command('echo Hello', None, a)
 136         assert a.getvalue() == 'Hello\n'
 137
 138         print "Reading from a stream to a StringIO..."
 139         file.seek(1)
 140         pipe_through_command('cat', file, a)
 141         assert a.getvalue() == 'Hello\nello World'
 142
 143         print "Writing to a fileno stream..."
 144         file.seek(0)
 145         file.truncate(0)
 146         pipe_through_command('echo Foo', None, file)
 147         file.seek(0)
 148         assert file.read() == 'Foo\n'
 149
 150         print "Read and write fileno streams..."
 151         src = Tmp()
 152         src.write('123')
 153         src.seek(0)
 154         file.seek(0)
 155         file.truncate(0)
 156         pipe_through_command('cat', src, file)
 157         file.seek(0)
 158         assert file.read() == '123'
 159
 160         print "Detect non-zero exit value..."
 161         try:
 162                 pipe_through_command('exit 1', None, None)
 163         except ChildError:
 164                 show()
 165         else:
 166                 assert 0
 167
 168         print "Detect writes to stderr..."
 169         try:
 170                 pipe_through_command('echo one >&2; sleep 2; echo two >&2', None, None)
 171         except ChildError:
 172                 show()
 173         else:
 174                 assert 0
 175
 176         print "Check tmp file is deleted..."
 177         name = file.name
 178         assert os.path.exists(name)
 179         file = None
 180         assert not os.path.exists(name)
 181
 182         print "All tests passed!"
 183
 184 if __name__ == '__main__':
 185         test()