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 _ValidateAndLogCommand(args, cwd, shell):
 119   if isinstance(args, basestring):
 120     if not shell:
 121       raise Exception('string args must be run with shell=True')
 122   else:
 123     if shell:
 124       raise Exception('array args must be run with shell=False')
 125     args = ' '.join(SingleQuote(c) for c in args)
 126   if cwd is None:
 127     cwd = ''
 128   else:
 129     cwd = ':' + cwd
 130   logging.info('[host]%s> %s', cwd, args)
 131   return args
 132
 133
 134 def GetCmdStatusAndOutput(args, cwd=None, shell=False):
 135   """Executes a subprocess and returns its exit code and output.
 136
 137   Args:
 138     args: A string or a sequence of program arguments. The program to execute is
 139       the string or the first item in the args sequence.
 140     cwd: If not None, the subprocess's current directory will be changed to
 141       |cwd| before it's executed.
 142     shell: Whether to execute args as a shell command. Must be True if args
 143       is a string and False if args is a sequence.
 144
 145   Returns:
 146     The 2-tuple (exit code, output).
 147   """
 148   _ValidateAndLogCommand(args, cwd, shell)
 149   pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
 150                shell=shell, cwd=cwd)
 151   stdout, stderr = pipe.communicate()
 152
 153   if stderr:
 154     logging.critical(stderr)
 155   if len(stdout) > 4096:
 156     logging.debug('Truncated output:')
 157   logging.debug(stdout[:4096])
 158   return (pipe.returncode, stdout)
 159
 160
 161 class TimeoutError(Exception):
 162   """Module-specific timeout exception."""
 163   pass
 164
 165
 166 def _IterProcessStdout(process, timeout=None, buffer_size=4096,
 167                        poll_interval=1):
 168   assert fcntl, 'fcntl module is required'
 169   try:
 170     # Enable non-blocking reads from the child's stdout.
 171     child_fd = process.stdout.fileno()
 172     fl = fcntl.fcntl(child_fd, fcntl.F_GETFL)
 173     fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
 174
 175     end_time = (time.time() + timeout) if timeout else None
 176     while True:
 177       if end_time and time.time() > end_time:
 178         raise TimeoutError
 179       read_fds, _, _ = select.select([child_fd], [], [], poll_interval)
 180       if child_fd in read_fds:
 181         data = os.read(child_fd, buffer_size)
 182         if not data:
 183           break
 184         yield data
 185       if process.poll() is not None:
 186         break
 187   finally:
 188     try:
 189       # Make sure the process doesn't stick around if we fail with an
 190       # exception.
 191       process.kill()
 192     except OSError:
 193       pass
 194     process.wait()
 195
 196
 197 def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False,
 198                                      logfile=None):
 199   """Executes a subprocess with a timeout.
 200
 201   Args:
 202     args: List of arguments to the program, the program to execute is the first
 203       element.
 204     timeout: the timeout in seconds or None to wait forever.
 205     cwd: If not None, the subprocess's current directory will be changed to
 206       |cwd| before it's executed.
 207     shell: Whether to execute args as a shell command. Must be True if args
 208       is a string and False if args is a sequence.
 209     logfile: Optional file-like object that will receive output from the
 210       command as it is running.
 211
 212   Returns:
 213     The 2-tuple (exit code, output).
 214   """
 215   _ValidateAndLogCommand(args, cwd, shell)
 216   output = StringIO.StringIO()
 217   process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE,
 218                   stderr=subprocess.STDOUT)
 219   for data in _IterProcessStdout(process, timeout=timeout):
 220     if logfile:
 221       logfile.write(data)
 222     output.write(data)
 223   return process.returncode, output.getvalue()
 224
 225
 226 def IterCmdOutputLines(args, timeout=None, cwd=None, shell=False,
 227                        check_status=True):
 228   """Executes a subprocess and continuously yields lines from its output.
 229
 230   Args:
 231     args: List of arguments to the program, the program to execute is the first
 232       element.
 233     cwd: If not None, the subprocess's current directory will be changed to
 234       |cwd| before it's executed.
 235     shell: Whether to execute args as a shell command. Must be True if args
 236       is a string and False if args is a sequence.
 237     check_status: A boolean indicating whether to check the exit status of the
 238       process after all output has been read.
 239
 240   Yields:
 241     The output of the subprocess, line by line.
 242
 243   Raises:
 244     CalledProcessError if check_status is True and the process exited with a
 245       non-zero exit status.
 246   """
 247   cmd = _ValidateAndLogCommand(args, cwd, shell)
 248   process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE,
 249                   stderr=subprocess.STDOUT)
 250   buffer_output = ''
 251   for data in _IterProcessStdout(process, timeout=timeout):
 252     buffer_output += data
 253     has_incomplete_line = buffer_output[-1] not in '\r\n'
 254     lines = buffer_output.splitlines()
 255     buffer_output = lines.pop() if has_incomplete_line else ''
 256     for line in lines:
 257       yield line
 258   if buffer_output:
 259     yield buffer_output
 260   if check_status and process.returncode:
 261     raise subprocess.CalledProcessError(process.returncode, cmd)