Home | History | Annotate | Download | only in pylib 
      1 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
      2 # Use of this source code is governed by a BSD-style license that can be
      3 # found in the LICENSE file.
      4 
      5 """A wrapper for subprocess to make calling shell commands easier."""
      6 
      7 import logging
      8 import os
      9 import pipes
     10 import select
     11 import signal
     12 import StringIO
     13 import subprocess
     14 import time
     15 
     16 # fcntl is not available on Windows.
     17 try:
     18   import fcntl
     19 except ImportError:
     20   fcntl = None
     21 
     22 
     23 def Popen(args, stdout=None, stderr=None, shell=None, cwd=None, env=None):
     24   return subprocess.Popen(
     25       args=args, cwd=cwd, stdout=stdout, stderr=stderr,
     26       shell=shell, close_fds=True, env=env,
     27       preexec_fn=lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL))
     28 
     29 
     30 def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None):
     31   pipe = Popen(args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd,
     32                env=env)
     33   pipe.communicate()
     34   return pipe.wait()
     35 
     36 
     37 def RunCmd(args, cwd=None):
     38   """Opens a subprocess to execute a program and returns its return value.
     39 
     40   Args:
     41     args: A string or a sequence of program arguments. The program to execute is
     42       the string or the first item in the args sequence.
     43     cwd: If not None, the subprocess's current directory will be changed to
     44       |cwd| before it's executed.
     45 
     46   Returns:
     47     Return code from the command execution.
     48   """
     49   logging.info(str(args) + ' ' + (cwd or ''))
     50   return Call(args, cwd=cwd)
     51 
     52 
     53 def GetCmdOutput(args, cwd=None, shell=False):
     54   """Open a subprocess to execute a program and returns its output.
     55 
     56   Args:
     57     args: A string or a sequence of program arguments. The program to execute is
     58       the string or the first item in the args sequence.
     59     cwd: If not None, the subprocess's current directory will be changed to
     60       |cwd| before it's executed.
     61     shell: Whether to execute args as a shell command.
     62 
     63   Returns:
     64     Captures and returns the command's stdout.
     65     Prints the command's stderr to logger (which defaults to stdout).
     66   """
     67   (_, output) = GetCmdStatusAndOutput(args, cwd, shell)
     68   return output
     69 
     70 
     71 def GetCmdStatusAndOutput(args, cwd=None, shell=False):
     72   """Executes a subprocess and returns its exit code and output.
     73 
     74   Args:
     75     args: A string or a sequence of program arguments. The program to execute is
     76       the string or the first item in the args sequence.
     77     cwd: If not None, the subprocess's current directory will be changed to
     78       |cwd| before it's executed.
     79     shell: Whether to execute args as a shell command.
     80 
     81   Returns:
     82     The 2-tuple (exit code, output).
     83   """
     84   if isinstance(args, basestring):
     85     args_repr = args
     86     if not shell:
     87       raise Exception('string args must be run with shell=True')
     88   elif shell:
     89     raise Exception('array args must be run with shell=False')
     90   else:
     91     args_repr = ' '.join(map(pipes.quote, args))
     92 
     93   s = '[host]'
     94   if cwd:
     95     s += ':' + cwd
     96   s += '> ' + args_repr
     97   logging.info(s)
     98   pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
     99                shell=shell, cwd=cwd)
    100   stdout, stderr = pipe.communicate()
    101 
    102   if stderr:
    103     logging.critical(stderr)
    104   if len(stdout) > 4096:
    105     logging.debug('Truncated output:')
    106   logging.debug(stdout[:4096])
    107   return (pipe.returncode, stdout)
    108 
    109 
    110 class TimeoutError(Exception):
    111   """Module-specific timeout exception."""
    112   pass
    113 
    114 
    115 def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False,
    116                                      logfile=None):
    117   """Executes a subprocess with a timeout.
    118 
    119   Args:
    120     args: List of arguments to the program, the program to execute is the first
    121       element.
    122     timeout: the timeout in seconds or None to wait forever.
    123     cwd: If not None, the subprocess's current directory will be changed to
    124       |cwd| before it's executed.
    125     shell: Whether to execute args as a shell command.
    126     logfile: Optional file-like object that will receive output from the
    127       command as it is running.
    128 
    129   Returns:
    130     The 2-tuple (exit code, output).
    131   """
    132   assert fcntl, 'fcntl module is required'
    133   process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE,
    134                   stderr=subprocess.STDOUT)
    135   try:
    136     end_time = (time.time() + timeout) if timeout else None
    137     poll_interval = 1
    138     buffer_size = 4096
    139     child_fd = process.stdout.fileno()
    140     output = StringIO.StringIO()
    141 
    142     # Enable non-blocking reads from the child's stdout.
    143     fl = fcntl.fcntl(child_fd, fcntl.F_GETFL)
    144     fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
    145 
    146     while True:
    147       if end_time and time.time() > end_time:
    148         raise TimeoutError
    149       read_fds, _, _ = select.select([child_fd], [], [], poll_interval)
    150       if child_fd in read_fds:
    151         data = os.read(child_fd, buffer_size)
    152         if not data:
    153           break
    154         if logfile:
    155           logfile.write(data)
    156         output.write(data)
    157       if process.poll() is not None:
    158         break
    159   finally:
    160     try:
    161       # Make sure the process doesn't stick around if we fail with an
    162       # exception.
    163       process.kill()
    164     except OSError:
    165       pass
    166     process.wait()
    167   return process.returncode, output.getvalue()
    168