Scott Zawalski | 6bc41ac | 2010-09-08 12:47:28 -0700 | [diff] [blame^] | 1 | # Copyright (c) 2010 The Chromium OS 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 | """Common python commands used by various build scripts.""" |
| 6 | |
| 7 | import subprocess |
| 8 | import sys |
| 9 | |
| 10 | _STDOUT_IS_TTY = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty() |
| 11 | |
| 12 | def RunCommand(cmd, print_cmd=True, error_ok=False, error_message=None, |
| 13 | exit_code=False, redirect_stdout=False, redirect_stderr=False, |
| 14 | cwd=None, input=None, enter_chroot=False): |
| 15 | """Runs a shell command. |
| 16 | |
| 17 | Keyword arguments: |
| 18 | cmd - cmd to run. Should be input to subprocess.POpen. If a string, |
| 19 | converted to an array using split(). |
| 20 | print_cmd -- prints the command before running it. |
| 21 | error_ok -- does not raise an exception on error. |
| 22 | error_message -- prints out this message when an error occurrs. |
| 23 | exit_code -- returns the return code of the shell command. |
| 24 | redirect_stdout -- returns the stdout. |
| 25 | redirect_stderr -- holds stderr output until input is communicated. |
| 26 | cwd -- the working directory to run this cmd. |
| 27 | input -- input to pipe into this command through stdin. |
| 28 | enter_chroot -- this command should be run from within the chroot. If set, |
| 29 | cwd must point to the scripts directory. |
| 30 | Raises: |
| 31 | Exception: Raises generic exception on error with optional error_message. |
| 32 | """ |
| 33 | # Set default for variables. |
| 34 | stdout = None |
| 35 | stderr = None |
| 36 | stdin = None |
| 37 | output = '' |
| 38 | |
| 39 | # Modify defaults based on parameters. |
| 40 | if redirect_stdout: stdout = subprocess.PIPE |
| 41 | if redirect_stderr: stderr = subprocess.PIPE |
| 42 | if input: stdin = subprocess.PIPE |
| 43 | if enter_chroot: cmd = ['./enter_chroot.sh', '--'] + cmd |
| 44 | |
| 45 | # Print out the command before running. |
| 46 | if print_cmd: |
| 47 | Info('RunCommand: %s' % ' '.join(cmd)) |
| 48 | |
| 49 | try: |
| 50 | proc = subprocess.Popen(cmd, cwd=cwd, stdin=stdin, |
| 51 | stdout=stdout, stderr=stderr) |
| 52 | (output, error) = proc.communicate(input) |
| 53 | if exit_code: |
| 54 | return proc.returncode |
| 55 | |
| 56 | if not error_ok and proc.returncode: |
| 57 | raise Exception('Command "%s" failed.\n' % (' '.join(cmd)) + |
| 58 | (error_message or error or output or '')) |
| 59 | except Exception,e: |
| 60 | if not error_ok: |
| 61 | raise |
| 62 | else: |
| 63 | Warning(str(e)) |
| 64 | |
| 65 | return output |
| 66 | |
| 67 | |
| 68 | class Color(object): |
| 69 | """Conditionally wraps text in ANSI color escape sequences.""" |
| 70 | BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8) |
| 71 | BOLD = -1 |
| 72 | COLOR_START = '\033[1;%dm' |
| 73 | BOLD_START = '\033[1m' |
| 74 | RESET = '\033[0m' |
| 75 | |
| 76 | def __init__(self, enabled=True): |
| 77 | self._enabled = enabled |
| 78 | |
| 79 | def Color(self, color, text): |
| 80 | """Returns text with conditionally added color escape sequences. |
| 81 | |
| 82 | Keyword arguments: |
| 83 | color: Text color -- one of the color constants defined in this class. |
| 84 | text: The text to color. |
| 85 | |
| 86 | Returns: |
| 87 | If self._enabled is False, returns the original text. If it's True, |
| 88 | returns text with color escape sequences based on the value of color. |
| 89 | """ |
| 90 | if not self._enabled: |
| 91 | return text |
| 92 | if color == self.BOLD: |
| 93 | start = self.BOLD_START |
| 94 | else: |
| 95 | start = self.COLOR_START % (color + 30) |
| 96 | return start + text + self.RESET |
| 97 | |
| 98 | |
| 99 | def Die(message): |
| 100 | """Emits a red error message and halts execution. |
| 101 | |
| 102 | Keyword arguments: |
| 103 | message: The message to be emitted before exiting. |
| 104 | """ |
| 105 | print >> sys.stderr, ( |
| 106 | Color(_STDOUT_IS_TTY).Color(Color.RED, '\nERROR: ' + message)) |
| 107 | sys.exit(1) |
| 108 | |
| 109 | |
| 110 | def Warning(message): |
| 111 | """Emits a yellow warning message and continues execution. |
| 112 | |
| 113 | Keyword arguments: |
| 114 | message: The message to be emitted. |
| 115 | """ |
| 116 | print >> sys.stderr, ( |
| 117 | Color(_STDOUT_IS_TTY).Color(Color.YELLOW, '\nWARNING: ' + message)) |
| 118 | |
| 119 | |
| 120 | def Info(message): |
| 121 | """Emits a blue informational message and continues execution. |
| 122 | |
| 123 | Keyword arguments: |
| 124 | message: The message to be emitted. |
| 125 | """ |
| 126 | print >> sys.stderr, ( |
| 127 | Color(_STDOUT_IS_TTY).Color(Color.BLUE, '\nINFO: ' + message)) |