build/android/pylib/cmd_helper.py

   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 string
  13 import StringIO
  14 import subprocess
  15 import time
  16
  17 # fcntl is not available on Windows.
  18 try:
  19   import fcntl
  20 except ImportError:
  21   fcntl = None
  22
  23 _SafeShellChars = frozenset(string.ascii_letters + string.digits + '@%_-+=:,./')
  24
  25 def SingleQuote(s):
  26   """Return an shell-escaped version of the string using single quotes.
  27
  28   Reliably quote a string which may contain unsafe characters (e.g. space,
  29   quote, or other special characters such as '$').
  30
  31   The returned value can be used in a shell command line as one token that gets
  32   to be interpreted literally.
  33
  34   Args:
  35     s: The string to quote.
  36
  37   Return:
  38     The string quoted using single quotes.
  39   """
  40   return pipes.quote(s)
  41
  42 def DoubleQuote(s):
  43   """Return an shell-escaped version of the string using double quotes.
  44
  45   Reliably quote a string which may contain unsafe characters (e.g. space
  46   or quote characters), while retaining some shell features such as variable
  47   interpolation.
  48
  49   The returned value can be used in a shell command line as one token that gets
  50   to be further interpreted by the shell.
  51
  52   The set of characters that retain their special meaning may depend on the
  53   shell implementation. This set usually includes: '$', '`', '\', '!', '*',
  54   and '@'.
  55
  56   Args:
  57     s: The string to quote.
  58
  59   Return:
  60     The string quoted using double quotes.
  61   """
  62   if not s:
  63     return '""'
  64   elif all(c in _SafeShellChars for c in s):
  65     return s
  66   else:
  67     return '"' + s.replace('"', '\\"') + '"'
  68
  69
  70 def Popen(args, stdout=None, stderr=None, shell=None, cwd=None, env=None):
  71   return subprocess.Popen(
  72       args=args, cwd=cwd, stdout=stdout, stderr=stderr,
  73       shell=shell, close_fds=True, env=env,
  74       preexec_fn=lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL))
  75
  76
  77 def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None):
  78   pipe = Popen(args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd,
  79                env=env)
  80   pipe.communicate()
  81   return pipe.wait()
  82
  83
  84 def RunCmd(args, cwd=None):
  85   """Opens a subprocess to execute a program and returns its return value.
  86
  87   Args:
  88     args: A string or a sequence of program arguments. The program to execute is
  89       the string or the first item in the args sequence.
  90     cwd: If not None, the subprocess's current directory will be changed to
  91       |cwd| before it's executed.
  92
  93   Returns:
  94     Return code from the command execution.
  95   """
  96   logging.info(str(args) + ' ' + (cwd or ''))
  97   return Call(args, cwd=cwd)
  98
  99
 100 def GetCmdOutput(args, cwd=None, shell=False):
 101   """Open a subprocess to execute a program and returns its output.
 102
 103   Args:
 104     args: A string or a sequence of program arguments. The program to execute is
 105       the string or the first item in the args sequence.
 106     cwd: If not None, the subprocess's current directory will be changed to
 107       |cwd| before it's executed.
 108     shell: Whether to execute args as a shell command.
 109
 110   Returns:
 111     Captures and returns the command's stdout.
 112     Prints the command's stderr to logger (which defaults to stdout).
 113   """
 114   (_, output) = GetCmdStatusAndOutput(args, cwd, shell)
 115   return output
 116
 117
 118 def _LogCommand(args, cwd=None):
 119   if not isinstance(args, basestring):
 120     args = ' '.join(SingleQuote(c) for c in args)
 121   if cwd is None:
 122     cwd = ''
 123   else:
 124     cwd = ':' + cwd
 125   logging.info('[host]%s> %s', cwd, args)
 126
 127
 128 def GetCmdStatusAndOutput(args, cwd=None, shell=False):
 129   """Executes a subprocess and returns its exit code and output.
 130
 131   Args:
 132     args: A string or a sequence of program arguments. The program to execute is
 133       the string or the first item in the args sequence.
 134     cwd: If not None, the subprocess's current directory will be changed to
 135       |cwd| before it's executed.
 136     shell: Whether to execute args as a shell command.
 137
 138   Returns:
 139     The 2-tuple (exit code, output).
 140   """
 141   if isinstance(args, basestring):
 142     if not shell:
 143       raise Exception('string args must be run with shell=True')
 144   elif shell:
 145     raise Exception('array args must be run with shell=False')
 146
 147   _LogCommand(args, cwd)
 148   pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
 149                shell=shell, cwd=cwd)
 150   stdout, stderr = pipe.communicate()
 151
 152   if stderr:
 153     logging.critical(stderr)
 154   if len(stdout) > 4096:
 155     logging.debug('Truncated output:')
 156   logging.debug(stdout[:4096])
 157   return (pipe.returncode, stdout)
 158
 159
 160 class TimeoutError(Exception):
 161   """Module-specific timeout exception."""
 162   pass
 163
 164
 165 def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False,
 166                                      logfile=None):
 167   """Executes a subprocess with a timeout.
 168
 169   Args:
 170     args: List of arguments to the program, the program to execute is the first
 171       element.
 172     timeout: the timeout in seconds or None to wait forever.
 173     cwd: If not None, the subprocess's current directory will be changed to
 174       |cwd| before it's executed.
 175     shell: Whether to execute args as a shell command.
 176     logfile: Optional file-like object that will receive output from the
 177       command as it is running.
 178
 179   Returns:
 180     The 2-tuple (exit code, output).
 181   """
 182   assert fcntl, 'fcntl module is required'
 183   if isinstance(args, basestring):
 184     if not shell:
 185       raise Exception('string args must be run with shell=True')
 186   elif shell:
 187     raise Exception('array args must be run with shell=False')
 188
 189   _LogCommand(args, cwd)
 190   process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE,
 191                   stderr=subprocess.STDOUT)
 192   try:
 193     end_time = (time.time() + timeout) if timeout else None
 194     poll_interval = 1
 195     buffer_size = 4096
 196     child_fd = process.stdout.fileno()
 197     output = StringIO.StringIO()
 198
 199     # Enable non-blocking reads from the child's stdout.
 200     fl = fcntl.fcntl(child_fd, fcntl.F_GETFL)
 201     fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
 202
 203     while True:
 204       if end_time and time.time() > end_time:
 205         raise TimeoutError
 206       read_fds, _, _ = select.select([child_fd], [], [], poll_interval)
 207       if child_fd in read_fds:
 208         data = os.read(child_fd, buffer_size)
 209         if not data:
 210           break
 211         if logfile:
 212           logfile.write(data)
 213         output.write(data)
 214       if process.poll() is not None:
 215         break
 216   finally:
 217     try:
 218       # Make sure the process doesn't stick around if we fail with an
 219       # exception.
 220       process.kill()
 221     except OSError:
 222       pass
 223     process.wait()
 224   return process.returncode, output.getvalue()