Copy subprocess2 from commit-queue so it can be reused here.

subprocess2.py

+# coding=utf8
+# Copyright (c) 2011 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.
+"""Collection of subprocess wrapper functions.
+
+In theory you shouldn't need anything else in subprocess, or this module failed.
+"""
+
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+import time
+import threading
+
+# Constants forwarded from subprocess.
+PIPE = subprocess.PIPE
+STDOUT = subprocess.STDOUT
+
+# Globals.
+# Set to True if you somehow need to disable this hack.
+SUBPROCESS_CLEANUP_HACKED = False
+
+
+class CalledProcessError(subprocess.CalledProcessError):
+  """Augment the standard exception with more data."""
+  def __init__(self, returncode, cmd, cwd, stdout, stderr):
+    super(CalledProcessError, self).__init__(returncode, cmd)
+    self.stdout = stdout
+    self.stderr = stderr
+    self.cwd = cwd
+
+  def __str__(self):
+    out = 'Command %s returned non-zero exit status %s' % (
+        ' '.join(self.cmd), self.returncode)
+    if self.cwd:
+      out += ' in ' + self.cwd
+    return '\n'.join(filter(None, (out, self.stdout, self.stderr)))
+
+
+def hack_subprocess():
+  """subprocess functions may throw exceptions when used in multiple threads.
+
+  See http://bugs.python.org/issue1731717 for more information.
+  """
+  global SUBPROCESS_CLEANUP_HACKED
+  if not SUBPROCESS_CLEANUP_HACKED and threading.activeCount() != 1:
+    # Only hack if there is ever multiple threads.
+    # There is no point to leak with only one thread.
+    subprocess._cleanup = lambda: None
+    SUBPROCESS_CLEANUP_HACKED = True
+
+
+def get_english_env(env):
+  """Forces LANG and/or LANGUAGE to be English.
+
+  Forces encoding to utf-8 for subprocesses.
+
+  Returns None if it is unnecessary.
+  """
+  env = env or os.environ
+
+  # Test if it is necessary at all.
+  is_english = lambda name: env.get(name, 'en').startswith('en')
+
+  if is_english('LANG') and is_english('LANGUAGE'):
+    return None
+
+  # Requires modifications.
+  env = env.copy()
+  def fix_lang(name):
+    if not is_english(name):
+      env[name] = 'en_US.UTF-8'
+  fix_lang('LANG')
+  fix_lang('LANGUAGE')
+  return env
+
+
+def Popen(args, **kwargs):
+  """Wraps subprocess.Popen().
+
+  Forces English output since it's easier to parse the stdout if it is always in
+  English.
+
+  Sets shell=True on windows by default. You can override this by forcing shell
+  parameter to a value.
+
+  Popen() can throw OSError when cwd or args[0] doesn't exist.
+  """
+  # Make sure we hack subprocess if necessary.
+  hack_subprocess()
+
+  env = get_english_env(kwargs.get('env'))
+  if env:
+    kwargs['env'] = env
+
+  if not kwargs.get('shell') is None:
+    # *Sigh*:  Windows needs shell=True, or else it won't search %PATH% for the
+    # executable, but shell=True makes subprocess on Linux fail when it's called
+    # with a list because it only tries to execute the first item in the list.
+    kwargs['shell'] = (sys.platform=='win32')
+
+  tmp_str = ' '.join(args)
+  if kwargs.get('cwd', None):
+    tmp_str += ';  cwd=%s' % kwargs['cwd']
+  logging.debug(tmp_str)
+  return subprocess.Popen(args, **kwargs)
+
+
+def call(args, timeout=None, **kwargs):
+  """Wraps subprocess.Popen().communicate().
+
+  The process will be kill with error code -9 after |timeout| seconds if set.
+
+  Automatically passes stdin content as input so do not specify stdin=PIPE.
+
+  Returns both communicate() tuple and return code wrapped in a tuple.
+  """
+  stdin = kwargs.pop('stdin', None)
+  if stdin is not None:
+    assert stdin != PIPE
+    # When stdin is passed as an argument, use it as the actual input data and
+    # set the Popen() parameter accordingly.
+    kwargs['stdin'] = PIPE
+
+  if not timeout:
+    # Normal workflow.
+    proc = Popen(args, **kwargs)
+    if stdin is not None:
+      out = proc.communicate(stdin)
+    else:
+      out = proc.communicate()
+  else:
+    # Create a temporary file to workaround python's deadlock.
+    # http://docs.python.org/library/subprocess.html#subprocess.Popen.wait
+    # When the pipe fills up, it will deadlock this process. Using a real file
+    # works around that issue.
+    with tempfile.TemporaryFile() as buff:
+      start = time.time()
+      kwargs['stdout'] = buff
+      proc = Popen(args, **kwargs)
+      if stdin is not None:
+        proc.stdin.write(stdin)
+      while proc.returncode is None:
+        proc.poll()
+        if timeout and (time.time() - start) > timeout:
+          proc.kill()
+          proc.wait()
+          # It's -9 on linux and 1 on Windows. Standardize to -9.
+          # Do not throw an exception here, the user must use
+          # check_call(timeout=60) and check for e.returncode == -9 instead.
+          # or look at call()[1] == -9.
+          proc.returncode = -9
+        time.sleep(0.001)
+      # Now that the process died, reset the cursor and read the file.
+      buff.seek(0)
+      out = [buff.read(), None]
+  return out, proc.returncode
+
+
+def check_call(args, **kwargs):
+  """Similar to subprocess.check_call() but use call() instead.
+
+  This permits to include more details in CalledProcessError().
+
+  Runs a command and throws an exception if the command failed.
+
+  Returns communicate() tuple.
+  """
+  out, returncode = call(args, **kwargs)
+  if returncode:
+    raise CalledProcessError(
+        returncode, args, kwargs.get('cwd'), out[0], out[1])
+  return out
+
+
+def capture(args, **kwargs):
+  """Captures stdout of a process call and returns it.
+
+  Similar to check_output() excepts that it discards return code.
+
+  Discards communicate()[1]. By default sets stderr=STDOUT.
+  """
+  if kwargs.get('stderr') is None:
+    kwargs['stderr'] = STDOUT
+  return call(args, stdout=PIPE, **kwargs)[0][0]
+
+
+def check_output(args, **kwargs):
+  """Captures stdout of a process call and returns it.
+
+  Discards communicate()[1]. By default sets stderr=STDOUT.
+
+  Throws if return code is not 0.
+
+  Works even prior to python 2.7.
+  """
+  if kwargs.get('stderr') is None:
+    kwargs['stderr'] = STDOUT
+  return check_call(args, stdout=PIPE, **kwargs)[0]