Add lots of utils, PRESUBMIT.py

py/utils/shell_utils.py

+#!/usr/bin/env python
+# Copyright (c) 2012 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+""" This module contains tools for running commands in a shell. """
+
+import datetime
+import os
+import Queue
+import select
+import subprocess
+import sys
+import threading
+import time
+
+if 'nt' in os.name:
+  import ctypes
+
+
+DEFAULT_SECS_BETWEEN_ATTEMPTS = 10
+POLL_MILLIS = 250
+
+
+class CommandFailedException(Exception):
+  """Exception which gets raised when a command fails."""
+
+  def __init__(self, output, *args):
+    """Initialize the CommandFailedException.
+
+    Args:
+        output: string; output from the command.
+    """
+    Exception.__init__(self, *args)
+    self._output = output
+
+  @property
+  def output(self):
+    """Output from the command."""
+    return self._output
+
+
+class TimeoutException(CommandFailedException):
+  """CommandFailedException which gets raised when a subprocess exceeds its
+  timeout. """
+  pass
+
+
+def run_async(cmd, echo=True, shell=False):
+  """ Run 'cmd' in a subprocess, returning a Popen class instance referring to
+  that process.  (Non-blocking) """
+  if echo:
+    print cmd
+  if 'nt' in os.name:
+    # Windows has a bad habit of opening a dialog when a console program
+    # crashes, rather than just letting it crash.  Therefore, when a program
+    # crashes on Windows, we don't find out until the build step times out.
+    # This code prevents the dialog from appearing, so that we find out
+    # immediately and don't waste time waiting around.
+    SEM_NOGPFAULTERRORBOX = 0x0002
+    ctypes.windll.kernel32.SetErrorMode(SEM_NOGPFAULTERRORBOX)
+    flags = 0x8000000 # CREATE_NO_WINDOW
+  else:
+    flags = 0
+  return subprocess.Popen(cmd, shell=shell, stderr=subprocess.STDOUT,
+                          stdout=subprocess.PIPE, creationflags=flags,
+                          bufsize=1)
+
+
+class EnqueueThread(threading.Thread):
+  """ Reads and enqueues lines from a file. """
+  def __init__(self, file_obj, queue):
+    threading.Thread.__init__(self)
+    self._file = file_obj
+    self._queue = queue
+    self._stopped = False
+
+  def run(self):
+    if sys.platform.startswith('linux'):
+      # Use a polling object to avoid the blocking call to readline().
+      poll = select.poll()
+      poll.register(self._file, select.POLLIN)
+      while not self._stopped:
+        has_output = poll.poll(POLL_MILLIS)
+        if has_output:
+          line = self._file.readline()
+          if line == '':
+            self._stopped = True
+          self._queue.put(line)
+    else:
+      # Only Unix supports polling objects, so just read from the file,
+      # Python-style.
+      for line in iter(self._file.readline, ''):
+        self._queue.put(line)
+        if self._stopped:
+          break
+
+  def stop(self):
+    self._stopped = True
+
+
+def log_process_in_real_time(proc, echo=True, timeout=None, log_file=None,
+                             halt_on_output=None, print_timestamps=True):
+  """ Log the output of proc in real time until it completes. Return a tuple
+  containing the exit code of proc and the contents of stdout.
+
+  proc: an instance of Popen referring to a running subprocess.
+  echo: boolean indicating whether to print the output received from proc.stdout
+  timeout: number of seconds allotted for the process to run. Raises a
+      TimeoutException if the run time exceeds the timeout.
+  log_file: an open file for writing output
+  halt_on_output: string; kill the process and return if this string is found
+      in the output stream from the process.
+  print_timestamps: boolean indicating whether a formatted timestamp should be
+      prepended to each line of output.
+  """
+  stdout_queue = Queue.Queue()
+  log_thread = EnqueueThread(proc.stdout, stdout_queue)
+  log_thread.start()
+  try:
+    all_output = []
+    t_0 = time.time()
+    while True:
+      code = proc.poll()
+      try:
+        output = stdout_queue.get_nowait()
+        all_output.append(output)
+        if output and print_timestamps:
+          timestamp = datetime.datetime.now().strftime('%H:%M:%S.%f')
+          output = ''.join(['[%s] %s\n' % (timestamp, line)
+                            for line in output.splitlines()])
+        if echo:
+          sys.stdout.write(output)
+          sys.stdout.flush()
+        if log_file:
+          log_file.write(output)
+          log_file.flush()
+        if halt_on_output and halt_on_output in output:
+          proc.terminate()
+          break
+      except Queue.Empty:
+        if code != None: # proc has finished running
+          break
+        time.sleep(0.5)
+      if timeout and time.time() - t_0 > timeout:
+        proc.terminate()
+        raise TimeoutException(
+            ''.join(all_output),
+            'Subprocess exceeded timeout of %ds' % timeout)
+  finally:
+    log_thread.stop()
+    log_thread.join()
+  return (code, ''.join(all_output))
+
+
+def log_process_after_completion(proc, echo=True, timeout=None, log_file=None):
+  """ Wait for proc to complete and return a tuple containing the exit code of
+  proc and the contents of stdout. Unlike log_process_in_real_time, does not
+  attempt to read stdout from proc in real time.
+
+  proc: an instance of Popen referring to a running subprocess.
+  echo: boolean indicating whether to print the output received from proc.stdout
+  timeout: number of seconds allotted for the process to run. Raises a
+      TimeoutException if the run time exceeds the timeout.
+  log_file: an open file for writing outout
+  """
+  t_0 = time.time()
+  code = None
+  while code is None:
+    if timeout and time.time() - t_0 > timeout:
+      raise TimeoutException(
+          proc.communicate()[0],
+          'Subprocess exceeded timeout of %ds' % timeout)
+    time.sleep(0.5)
+    code = proc.poll()
+  output = proc.communicate()[0]
+  if echo:
+    print output
+  if log_file:
+    log_file.write(output)
+    log_file.flush()
+  return (code, output)
+
+
+def run(cmd, echo=True, shell=False, timeout=None, print_timestamps=True,
+        log_in_real_time=True):
+  """ Run 'cmd' in a shell and return the combined contents of stdout and
+  stderr (Blocking).  Throws an exception if the command exits non-zero.
+
+  cmd: list of strings (or single string, iff shell==True) indicating the
+      command to run
+  echo: boolean indicating whether we should print the command and log output
+  shell: boolean indicating whether we are using advanced shell features. Use
+      only when absolutely necessary, since this allows a lot more freedom which
+      could be exploited by malicious code. See the warning here:
+      http://docs.python.org/library/subprocess.html#popen-constructor
+  timeout: optional, integer indicating the maximum elapsed time in seconds
+  print_timestamps: boolean indicating whether a formatted timestamp should be
+      prepended to each line of output. Unused if echo or log_in_real_time is
+      False.
+  log_in_real_time: boolean indicating whether to read stdout from the
+      subprocess in real time instead of when the process finishes. If echo is
+      False, we never log in real time, even if log_in_real_time is True.
+  """
+  proc = run_async(cmd, echo=echo, shell=shell)
+  # If we're not printing the output, we don't care if the output shows up in
+  # real time, so don't bother.
+  if log_in_real_time and echo:
+    (returncode, output) = log_process_in_real_time(proc, echo=echo,
+        timeout=timeout, print_timestamps=print_timestamps)
+  else:
+    (returncode, output) = log_process_after_completion(proc, echo=echo,
+                                                        timeout=timeout)
+  if returncode != 0:
+    raise CommandFailedException(
+        output,
+        'Command failed with code %d: %s' % (returncode, cmd))
+  return output
+
+
+def run_retry(cmd, echo=True, shell=False, attempts=1,
+              secs_between_attempts=DEFAULT_SECS_BETWEEN_ATTEMPTS,
+              timeout=None, print_timestamps=True):
+  """ Wrapper for run() which makes multiple attempts until either the command
+  succeeds or the maximum number of attempts is reached. """
+  attempt = 1
+  while True:
+    try:
+      return run(cmd, echo=echo, shell=shell, timeout=timeout,
+                 print_timestamps=print_timestamps)
+    except CommandFailedException:
+      if attempt >= attempts:
+        raise
+    print 'Command failed. Retrying in %d seconds...' % secs_between_attempts
+    time.sleep(secs_between_attempts)
+    attempt += 1