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.
"""

from __future__ import with_statement
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]