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 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()