Move RunCommand, and Info/Warning/Die into common pylib

TEST=Ran all commands that used RunCommand or Die with
their associated unittests.

Review URL: http://codereview.chromium.org/3266004
This commit is contained in:
Chris Sosa 2010-08-30 11:05:50 -07:00
parent 1ec70345d9
commit 70e2c9e72a
7 changed files with 204 additions and 166 deletions

View File

@ -101,6 +101,10 @@ function main() {
mv autotest "tarball/${FLAGS_output_tag}"
cp "${script_dir}/generate_test_report.py" \
"tarball/${FLAGS_output_tag}/generate_test_report"
# Copy python lib used in generate_test_report.
mkdir -p "tarball/${FLAGS_output_tag}/lib"
cp "${script_dir}/lib/cros_build_lib.py" \
"tarball/${FLAGS_output_tag}/lib"
echo "Creating ${FLAGS_to}/${FLAGS_output_tag}.tar.bz2..."
mkdir -p "${FLAGS_to}"
cd tarball

View File

@ -10,61 +10,18 @@ import errno
import re
import optparse
import os
import subprocess
import sys
import cbuildbot_comm
from cbuildbot_config import config
sys.path.append(os.path.join(os.path.dirname(__file__), '../lib'))
from cros_build_lib import Die, Info, RunCommand, Warning
_DEFAULT_RETRIES = 3
# ======================== Utility functions ================================
def RunCommand(cmd, print_cmd=True, error_ok=False, error_message=None,
exit_code=False, redirect_stdout=False, redirect_stderr=False,
cwd=None, input=None, enter_chroot=False):
"""Runs a shell command.
Keyword arguments:
print_cmd -- prints the command before running it.
error_ok -- does not raise an exception on error.
error_message -- prints out this message when an error occurrs.
exit_code -- returns the return code of the shell command.
redirect_stdout -- returns the stdout.
redirect_stderr -- holds stderr output until input is communicated.
cwd -- the working directory to run this cmd.
input -- input to pipe into this command through stdin.
enter_chroot -- this command should be run from within the chroot.
"""
# Set default for variables.
stdout = None
stderr = None
stdin = None
# Modify defaults based on parameters.
if redirect_stdout: stdout = subprocess.PIPE
if redirect_stderr: stderr = subprocess.PIPE
if input: stdin = subprocess.PIPE
if enter_chroot: cmd = ['./enter_chroot.sh', '--'] + cmd
# Print out the command before running.
if print_cmd:
print >> sys.stderr, 'CBUILDBOT -- RunCommand: ', ' '.join(cmd)
proc = subprocess.Popen(cmd, cwd=cwd, stdin=stdin,
stdout=stdout, stderr=stderr)
(output, error) = proc.communicate(input)
if exit_code:
return proc.returncode
if not error_ok and proc.returncode != 0:
raise Exception('Command "%s" failed.\n' % (' '.join(cmd)) +
(error_message or error or output or ''))
return output
def MakeDir(path, parents=False):
"""Basic wrapper around os.mkdirs.
@ -104,9 +61,9 @@ def RepoSync(buildroot, rw_checkout=False, retries=_DEFAULT_RETRIES):
except:
retries -= 1
if retries > 0:
print >> sys.stderr, 'CBUILDBOT -- Repo Sync Failed, retrying'
Warning('CBUILDBOT -- Repo Sync Failed, retrying')
else:
print >> sys.stderr, 'CBUILDBOT -- Retries exhausted'
Warning('CBUILDBOT -- Retries exhausted')
raise
# =========================== Command Helpers =================================
@ -129,7 +86,7 @@ def _GetAllGitRepos(buildroot, debug=False):
# Create the array.
for result in result_array:
if len(result) != 2:
print >> sys.stderr, 'Found in correct xml object %s', result
Warning('Found incorrect xml object %s' % result)
else:
# Remove pre-pended src directory from manifest.
manifest_tuples.append([result[0], result[1].replace('src/', '')])
@ -161,8 +118,7 @@ def _CreateRepoDictionary(buildroot, board, debug=False):
"""Returns the repo->list_of_ebuilds dictionary."""
repo_dictionary = {}
manifest_tuples = _GetAllGitRepos(buildroot)
print >> sys.stderr, (
'Creating dictionary of git repos to portage packages ...')
Info('Creating dictionary of git repos to portage packages ...')
cwd = os.path.join(buildroot, 'src', 'scripts')
get_all_workon_pkgs_cmd = './cros_workon list --all'.split()
@ -175,8 +131,7 @@ def _CreateRepoDictionary(buildroot, board, debug=False):
for tuple in manifest_tuples:
# This path tends to have the user's home_dir prepended to it.
if cros_workon_src_path.endswith(tuple[1]):
print >> sys.stderr, ('For %s found matching package %s' %
(tuple[0], package))
Info('For %s found matching package %s' % (tuple[0], package))
if repo_dictionary.has_key(tuple[0]):
repo_dictionary[tuple[0]] += [package]
else:
@ -204,7 +159,7 @@ def _ParseRevisionString(revision_string, repo_dictionary):
# Format 'package@commit-id'.
revision_tuple = revision.split('@')
if len(revision_tuple) != 2:
print >> sys.stderr, 'Incorrectly formatted revision %s' % revision
Warning('Incorrectly formatted revision %s' % revision)
repo_name = revision_tuple[0].replace('.git', '')
# Might not have entry if no matching ebuild.
@ -219,7 +174,7 @@ def _ParseRevisionString(revision_string, repo_dictionary):
def _UprevFromRevisionList(buildroot, revision_list):
"""Uprevs based on revision list."""
if not revision_list:
print >> sys.stderr, 'No packages found to uprev'
Info('No packages found to uprev')
return
package_str = ''
@ -309,21 +264,20 @@ def _UprevPackages(buildroot, revisionfile, board):
revisions = rev_file.read()
rev_file.close()
except Exception, e:
print >> sys.stderr, 'Error reading %s, revving all' % revisionfile
print e
Warning('Error reading %s, revving all' % revisionfile)
revisions = 'None'
revisions = revisions.strip()
# TODO(sosa): Un-comment once we close individual trees.
# Revisions == "None" indicates a Force Build.
# revisions == "None" indicates a Force Build.
#if revisions != 'None':
# print >> sys.stderr, 'CBUILDBOT Revision list found %s' % revisions
# revision_list = _ParseRevisionString(revisions,
# _CreateRepoDictionary(buildroot, board))
# _UprevFromRevisionList(buildroot, revision_list)
#else:
print >> sys.stderr, 'CBUILDBOT Revving all'
Info('CBUILDBOT Revving all')
_UprevAllPackages(buildroot)
@ -349,12 +303,12 @@ def _GetConfig(config_name):
default = config['default']
buildconfig = {}
if not config.has_key(config_name):
print >> sys.stderr, 'Non-existent configuration specified.'
print >> sys.stderr, 'Please specify one of:'
Warning('Non-existent configuration specified.')
Warning('Please specify one of:')
config_names = config.keys()
config_names.sort()
for name in config_names:
print >> sys.stderr, ' %s' % name
Warning(' %s' % name)
sys.exit(1)
buildconfig = config[config_name]
@ -391,7 +345,7 @@ def main():
if len(args) == 1:
buildconfig = _GetConfig(args[0])
else:
print >> sys.stderr, "Missing configuration description"
Warning('Missing configuration description')
parser.print_usage()
sys.exit(1)
@ -426,8 +380,7 @@ def main():
else:
# At least one of the slaves failed or we timed out.
_UprevCleanup(buildroot)
sys.stderr('CBUILDBOT - One of the slaves has failed!!!')
sys.exit(1)
Die('CBUILDBOT - One of the slaves has failed!!!')
else:
# Publish my status to the master if its expecting it.
if buildconfig['important']:

View File

@ -6,11 +6,13 @@
import Queue
import SocketServer
import os
import socket
import sys
import time
from cbuildbot import RunCommand
sys.path.append(os.path.join(os.path.dirname(__file__), '../lib'))
from cros_build_lib import Info, Warning, RunCommand
# Communication port for master to slave communication.
_COMM_PORT = 32890
@ -52,15 +54,14 @@ class _SlaveCommandHandler(SocketServer.BaseRequestHandler):
def _HandleCommand(self, command, args):
"""Handles command and returns status for master."""
print >> sys.stderr, ('(Slave) - Received command %s with args %s' %
(command, args))
Info('(Slave) - Received command %s with args %s' % (command, args))
command_to_expect = _command_queue.get()
# Check status also adds an entry on the status queue.
if command_to_expect == _COMMAND_CHECK_STATUS:
slave_status = _status_queue.get()
# Safety check to make sure the server is in a good state.
if command_to_expect != command:
print >> sys.stderr, (
Warning(
'(Slave) - Rejecting command %s. Was expecting %s.' % (command,
command_to_expect))
return _STATUS_COMMAND_REJECTED
@ -91,7 +92,7 @@ def _GetSlaveNames(configuration):
def _SendCommand(hostname, command, args):
"""Returns response from host or _STATUS_TIMEOUT on error."""
data = '%s\n%s\n' % (command, args)
print '(Master) - Sending %s %s to %s' % (command, args, hostname)
Info('(Master) - Sending %s %s to %s' % (command, args, hostname))
# Create a socket (SOCK_STREAM means a TCP socket).
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@ -124,10 +125,10 @@ def _CheckSlavesLeftStatus(slaves_to_check):
for slave in slaves_to_check:
status = _SendCommand(slave, _COMMAND_CHECK_STATUS, 'empty')
if status == STATUS_BUILD_FAILED:
print >> sys.stderr, '(Master) - Slave %s failed' % slave
Warning('(Master) - Slave %s failed' % slave)
return False
elif status == STATUS_BUILD_COMPLETE:
print >> sys.stderr, '(Master) - Slave %s completed' % slave
Info('(Master) - Slave %s completed' % slave)
slaves_to_remove.append(slave)
for slave in slaves_to_remove:
slaves_to_check.remove(slave)
@ -185,11 +186,10 @@ def PublishStatus(status):
try:
response = _receive_queue.get_nowait()
except Queue.Empty:
print >> sys.stderr, ('(Slave) - Waiting for master to accept %s' % (
status))
Info('(Slave) - Waiting for master to accept %s' % status)
timeout += _HEARTBEAT_TIMEOUT
response = None
except Exception, e:
print >> sys.stderr, '%s' % e
Warning('%s' % e)
server.server_close()
return response != None

View File

@ -15,9 +15,8 @@ import shutil
import subprocess
import sys
# TODO(sosa): Refactor Die into common library.
sys.path.append(os.path.dirname(__file__))
import generate_test_report
sys.path.append(os.path.join(os.path.dirname(__file__), 'lib'))
from cros_build_lib import Info, Warning, Die
gflags.DEFINE_string('board', 'x86-generic',
@ -33,7 +32,7 @@ gflags.DEFINE_string('packages', '',
short_name='p')
gflags.DEFINE_string('push_options', '',
'Options to use with git-cl push using push command.')
gflags.DEFINE_string('srcroot', '%s/trunk/src' % os.environ['HOME'],
gflags.DEFINE_string('srcroot', '%s/trunk/src' % os.environ['HOME'],
'Path to root src directory.',
short_name='r')
gflags.DEFINE_string('tracking_branch', 'cros/master',
@ -67,11 +66,11 @@ _STABLE_BRANCH_NAME = 'stabilizing_branch'
def _Print(message):
"""Verbose print function."""
if gflags.FLAGS.verbose:
print message
Info(message)
def _CheckOnStabilizingBranch():
"""Returns true if the git branch is on the stabilizing branch."""
current_branch = _RunCommand('git branch | grep \*').split()[1]
current_branch = _SimpleRunCommand('git branch | grep \*').split()[1]
return current_branch == _STABLE_BRANCH_NAME
def _CheckSaneArguments(package_list, commit_id_list, command):
@ -91,8 +90,8 @@ def _CheckSaneArguments(package_list, commit_id_list, command):
def _Clean():
"""Cleans up uncommitted changes on either stabilizing branch or master."""
_RunCommand('git reset HEAD --hard')
_RunCommand('git checkout %s' % gflags.FLAGS.tracking_branch)
_SimpleRunCommand('git reset HEAD --hard')
_SimpleRunCommand('git checkout %s' % gflags.FLAGS.tracking_branch)
def _PrintUsageAndDie(error_message=''):
@ -103,10 +102,10 @@ def _PrintUsageAndDie(error_message=''):
for command in commands:
command_usage += ' %s: %s\n' % (command, _COMMAND_DICTIONARY[command])
commands_str = '|'.join(commands)
print 'Usage: %s FLAGS [%s]\n\n%s\nFlags:%s' % (sys.argv[0], commands_str,
command_usage, gflags.FLAGS)
Warning('Usage: %s FLAGS [%s]\n\n%s\nFlags:%s' % (sys.argv[0], commands_str,
command_usage, gflags.FLAGS))
if error_message:
generate_test_report.Die(error_message)
Die(error_message)
else:
sys.exit(1)
@ -125,27 +124,27 @@ def _PushChange():
# Sanity check to make sure we're on a stabilizing branch before pushing.
if not _CheckOnStabilizingBranch():
print 'Not on branch %s so no work found to push. Exiting' % \
_STABLE_BRANCH_NAME
Info('Not on branch %s so no work found to push. Exiting' % \
_STABLE_BRANCH_NAME)
return
description = _RunCommand('git log --format=format:%s%n%n%b ' +
description = _SimpleRunCommand('git log --format=format:%s%n%n%b ' +
gflags.FLAGS.tracking_branch + '..')
description = 'Marking set of ebuilds as stable\n\n%s' % description
merge_branch_name = 'merge_branch'
_RunCommand('git remote update')
_SimpleRunCommand('git remote update')
merge_branch = _GitBranch(merge_branch_name)
merge_branch.CreateBranch()
if not merge_branch.Exists():
generate_test_report.Die('Unable to create merge branch.')
_RunCommand('git merge --squash %s' % _STABLE_BRANCH_NAME)
_RunCommand('git commit -m "%s"' % description)
Die('Unable to create merge branch.')
_SimpleRunCommand('git merge --squash %s' % _STABLE_BRANCH_NAME)
_SimpleRunCommand('git commit -m "%s"' % description)
# Ugh. There has got to be an easier way to push to a tracking branch
_RunCommand('git config push.default tracking')
_RunCommand('git push')
_SimpleRunCommand('git config push.default tracking')
_SimpleRunCommand('git push')
def _RunCommand(command):
def _SimpleRunCommand(command):
"""Runs a shell command and returns stdout back to caller."""
_Print(' + %s' % command)
proc_handle = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)
@ -174,12 +173,12 @@ class _GitBranch(object):
git_cmd = 'git checkout -b %s %s' % (target, gflags.FLAGS.tracking_branch)
else:
git_cmd = 'git checkout %s' % target
_RunCommand(git_cmd)
_SimpleRunCommand(git_cmd)
def Exists(self):
"""Returns True if the branch exists."""
branch_cmd = 'git branch'
branches = _RunCommand(branch_cmd)
branches = _SimpleRunCommand(branch_cmd)
return self.branch_name in branches.split()
def Delete(self):
@ -189,7 +188,7 @@ class _GitBranch(object):
"""
self._Checkout(gflags.FLAGS.tracking_branch, create=False)
delete_cmd = 'git branch -D %s' % self.branch_name
_RunCommand(delete_cmd)
_SimpleRunCommand(delete_cmd)
class _EBuild(object):
@ -214,7 +213,7 @@ class _EBuild(object):
_Print('Looking for unstable ebuild for %s' % package)
equery_cmd = 'equery-%s which %s 2> /dev/null' \
% (gflags.FLAGS.board, package)
path = _RunCommand(equery_cmd)
path = _SimpleRunCommand(equery_cmd)
if path:
_Print('Unstable ebuild found at %s' % path)
return path
@ -270,7 +269,7 @@ class EBuildStableMarker(object):
"""
# TODO(sosa): Change to a check.
if not self._ebuild:
generate_test_report.Die('Invalid ebuild given to EBuildStableMarker')
Die('Invalid ebuild given to EBuildStableMarker')
new_ebuild_path = '%s-r%d.ebuild' % (self._ebuild.ebuild_path_no_revision,
self._ebuild.current_revision + 1)
@ -296,10 +295,10 @@ class EBuildStableMarker(object):
fileinput.close()
_Print('Adding new stable ebuild to git')
_RunCommand('git add %s' % new_ebuild_path)
_SimpleRunCommand('git add %s' % new_ebuild_path)
_Print('Removing old ebuild from git')
_RunCommand('git rm %s' % self._ebuild.ebuild_path)
_SimpleRunCommand('git rm %s' % self._ebuild.ebuild_path)
def CommitChange(self, message):
"""Commits current changes in git locally.
@ -316,7 +315,7 @@ class EBuildStableMarker(object):
_Print('Committing changes for %s with commit message %s' % \
(self._ebuild.package, message))
git_commit_cmd = 'git commit -am "%s"' % message
_RunCommand(git_commit_cmd)
_SimpleRunCommand(git_commit_cmd)
def main(argv):
@ -346,7 +345,7 @@ def main(argv):
work_branch = _GitBranch(_STABLE_BRANCH_NAME)
work_branch.CreateBranch()
if not work_branch.Exists():
generate_test_report.Die('Unable to create stabilizing branch in %s' %
Die('Unable to create stabilizing branch in %s' %
overlay_directory)
index = 0
try:
@ -363,11 +362,11 @@ def main(argv):
worker.CommitChange(_GIT_COMMIT_MESSAGE % (package, commit_id))
except (OSError, IOError), e:
print ('An exception occurred\n'
'Only the following packages were revved: %s\n'
'Note you will have to go into %s'
'and reset the git repo yourself.' %
(package_list[:index], overlay_directory))
Warning('An exception occurred\n'
'Only the following packages were revved: %s\n'
'Note you will have to go into %s'
'and reset the git repo yourself.' %
(package_list[:index], overlay_directory))
raise e
elif command == 'push':
_PushChange()

View File

@ -21,7 +21,7 @@ class GitBranchTest(mox.MoxTestBase):
def setUp(self):
mox.MoxTestBase.setUp(self)
# Always stub RunCommmand out as we use it in every method.
self.mox.StubOutWithMock(cros_mark_as_stable, '_RunCommand')
self.mox.StubOutWithMock(cros_mark_as_stable, '_SimpleRunCommand')
self._branch = 'test_branch'
def testCreateBranchNoPrevious(self):
@ -50,7 +50,8 @@ class GitBranchTest(mox.MoxTestBase):
def testCheckoutCreate(self):
# Test init with no previous branch existing.
cros_mark_as_stable._RunCommand('git checkout -b %s origin' % self._branch)
cros_mark_as_stable._SimpleRunCommand(
'git checkout -b %s cros/master' % self._branch)
self.mox.ReplayAll()
branch = cros_mark_as_stable._GitBranch(self._branch)
branch._Checkout(self._branch)
@ -58,17 +59,17 @@ class GitBranchTest(mox.MoxTestBase):
def testCheckoutNoCreate(self):
# Test init with previous branch existing.
cros_mark_as_stable._RunCommand('git checkout master')
cros_mark_as_stable._SimpleRunCommand('git checkout cros/master')
self.mox.ReplayAll()
branch = cros_mark_as_stable._GitBranch(self._branch)
branch._Checkout('master', False)
branch._Checkout('cros/master', False)
self.mox.VerifyAll()
def testDelete(self):
branch = cros_mark_as_stable._GitBranch(self._branch)
self.mox.StubOutWithMock(branch, '_Checkout')
branch._Checkout('master', create=False)
cros_mark_as_stable._RunCommand('git branch -D ' + self._branch)
branch._Checkout('cros/master', create=False)
cros_mark_as_stable._SimpleRunCommand('git branch -D ' + self._branch)
self.mox.ReplayAll()
branch.Delete()
self.mox.VerifyAll()
@ -77,8 +78,8 @@ class GitBranchTest(mox.MoxTestBase):
branch = cros_mark_as_stable._GitBranch(self._branch)
# Test if branch exists that is created
cros_mark_as_stable._RunCommand('git branch').AndReturn(
'%s %s' % (self._branch, 'master'))
cros_mark_as_stable._SimpleRunCommand('git branch').AndReturn(
'%s %s' % (self._branch, 'cros/master'))
self.mox.ReplayAll()
self.assertTrue(branch.Exists())
self.mox.VerifyAll()
@ -114,8 +115,8 @@ class EBuildTest(mox.MoxTestBase):
self.assertEquals(ebuild.commit_id, 'my_id')
def testFindEBuildPath(self):
self.mox.StubOutWithMock(cros_mark_as_stable, '_RunCommand')
cros_mark_as_stable._RunCommand(
self.mox.StubOutWithMock(cros_mark_as_stable, '_SimpleRunCommand')
cros_mark_as_stable._SimpleRunCommand(
'equery-x86-generic which %s 2> /dev/null' % self.package).AndReturn(
self.ebuild_path)
self.mox.ReplayAll()
@ -144,7 +145,7 @@ class EBuildStableMarkerTest(mox.MoxTestBase):
def setUp(self):
mox.MoxTestBase.setUp(self)
self.mox.StubOutWithMock(cros_mark_as_stable, '_RunCommand')
self.mox.StubOutWithMock(cros_mark_as_stable, '_SimpleRunCommand')
self.m_ebuild = self.mox.CreateMock(cros_mark_as_stable._EBuild)
self.m_ebuild.package = 'test_package'
self.m_ebuild.current_revision = 1
@ -172,8 +173,8 @@ class EBuildStableMarkerTest(mox.MoxTestBase):
m_file.write('CROS_WORKON_COMMIT="my_id"\n')
m_file.write('KEYWORDS="x86 arm"')
m_file.write('src_unpack(){}')
cros_mark_as_stable._RunCommand('git add ' + self.revved_ebuild_path)
cros_mark_as_stable._RunCommand('git rm ' + self.m_ebuild.ebuild_path)
cros_mark_as_stable._SimpleRunCommand('git add ' + self.revved_ebuild_path)
cros_mark_as_stable._SimpleRunCommand('git rm ' + self.m_ebuild.ebuild_path)
self.mox.ReplayAll()
marker = cros_mark_as_stable.EBuildStableMarker(self.m_ebuild)
@ -183,7 +184,7 @@ class EBuildStableMarkerTest(mox.MoxTestBase):
def testCommitChange(self):
mock_message = 'Commit me'
cros_mark_as_stable._RunCommand(
cros_mark_as_stable._SimpleRunCommand(
'git commit -am "%s"' % mock_message)
self.mox.ReplayAll()
marker = cros_mark_as_stable.EBuildStableMarker(self.m_ebuild)
@ -191,7 +192,7 @@ class EBuildStableMarkerTest(mox.MoxTestBase):
self.mox.VerifyAll()
def testPushChange(self):
#cros_mark_as_stable._RunCommand('git push')
#cros_mark_as_stable._SimpleRunCommand('git push')
#self.mox.ReplayAll()
#marker = cros_mark_as_stable.EBuildStableMarker(self.m_ebuild)
#marker.PushChange()

View File

@ -17,51 +17,12 @@ import os
import re
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), 'lib'))
from cros_build_lib import Color, Die
_STDOUT_IS_TTY = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()
class Color(object):
"""Conditionally wraps text in ANSI color escape sequences."""
BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
BOLD = -1
COLOR_START = '\033[1;%dm'
BOLD_START = '\033[1m'
RESET = '\033[0m'
def __init__(self, enabled=True):
self._enabled = enabled
def Color(self, color, text):
"""Returns text with conditionally added color escape sequences.
Args:
color: Text color -- one of the color constants defined in this class.
text: The text to color.
Returns:
If self._enabled is False, returns the original text. If it's True,
returns text with color escape sequences based on the value of color.
"""
if not self._enabled:
return text
if color == self.BOLD:
start = self.BOLD_START
else:
start = self.COLOR_START % (color + 30)
return start + text + self.RESET
def Die(message):
"""Emits a red error message and halts execution.
Args:
message: The message to be emitted before exiting.
"""
print Color(_STDOUT_IS_TTY).Color(Color.RED, '\nERROR: ' + message)
sys.exit(1)
class ReportGenerator(object):
"""Collects and displays data from autoserv results directories.

120
lib/cros_build_lib.py Normal file
View File

@ -0,0 +1,120 @@
# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Common python commands used by various build scripts."""
import subprocess
import sys
_STDOUT_IS_TTY = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()
def RunCommand(cmd, print_cmd=True, error_ok=False, error_message=None,
exit_code=False, redirect_stdout=False, redirect_stderr=False,
cwd=None, input=None, enter_chroot=False):
"""Runs a shell command.
Keyword arguments:
cmd - cmd to run. Should be input to subprocess.POpen. If a string,
converted to an array using split().
print_cmd -- prints the command before running it.
error_ok -- does not raise an exception on error.
error_message -- prints out this message when an error occurrs.
exit_code -- returns the return code of the shell command.
redirect_stdout -- returns the stdout.
redirect_stderr -- holds stderr output until input is communicated.
cwd -- the working directory to run this cmd.
input -- input to pipe into this command through stdin.
enter_chroot -- this command should be run from within the chroot. If set,
cwd must point to the scripts directory.
Raises:
Exception: Raises generic exception on error with optional error_message.
"""
# Set default for variables.
stdout = None
stderr = None
stdin = None
# Modify defaults based on parameters.
if redirect_stdout: stdout = subprocess.PIPE
if redirect_stderr: stderr = subprocess.PIPE
if input: stdin = subprocess.PIPE
if enter_chroot: cmd = ['./enter_chroot.sh', '--'] + cmd
# Print out the command before running.
if print_cmd:
Info('RunCommand: %s' % ' '.join(cmd))
proc = subprocess.Popen(cmd, cwd=cwd, stdin=stdin,
stdout=stdout, stderr=stderr)
(output, error) = proc.communicate(input)
if exit_code:
return proc.returncode
if not error_ok and proc.returncode:
raise Exception('Command "%s" failed.\n' % (' '.join(cmd)) +
(error_message or error or output or ''))
return output
class Color(object):
"""Conditionally wraps text in ANSI color escape sequences."""
BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
BOLD = -1
COLOR_START = '\033[1;%dm'
BOLD_START = '\033[1m'
RESET = '\033[0m'
def __init__(self, enabled=True):
self._enabled = enabled
def Color(self, color, text):
"""Returns text with conditionally added color escape sequences.
Keyword arguments:
color: Text color -- one of the color constants defined in this class.
text: The text to color.
Returns:
If self._enabled is False, returns the original text. If it's True,
returns text with color escape sequences based on the value of color.
"""
if not self._enabled:
return text
if color == self.BOLD:
start = self.BOLD_START
else:
start = self.COLOR_START % (color + 30)
return start + text + self.RESET
def Die(message):
"""Emits a red error message and halts execution.
Keyword arguments:
message: The message to be emitted before exiting.
"""
print >> sys.stderr, (
Color(_STDOUT_IS_TTY).Color(Color.RED, '\nERROR: ' + message))
sys.exit(1)
def Warning(message):
"""Emits a yellow warning message and continues execution.
Keyword arguments:
message: The message to be emitted.
"""
print >> sys.stderr, (
Color(_STDOUT_IS_TTY).Color(Color.YELLOW, '\nWARNING: ' + message))
def Info(message):
"""Emits a blue informational message and continues execution.
Keyword arguments:
message: The message to be emitted.
"""
print >> sys.stderr, (
Color(_STDOUT_IS_TTY).Color(Color.BLUE, '\nINFO: ' + message))