| OLD | NEW |
| (Empty) |
| 1 # Copyright (c) 2012 The Chromium 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 """A wrapper for subprocess to make calling shell commands easier.""" | |
| 6 | |
| 7 import logging | |
| 8 import os | |
| 9 import pipes | |
| 10 import select | |
| 11 import signal | |
| 12 import string | |
| 13 import StringIO | |
| 14 import subprocess | |
| 15 import time | |
| 16 | |
| 17 # fcntl is not available on Windows. | |
| 18 try: | |
| 19 import fcntl | |
| 20 except ImportError: | |
| 21 fcntl = None | |
| 22 | |
| 23 _SafeShellChars = frozenset(string.ascii_letters + string.digits + '@%_-+=:,./') | |
| 24 | |
| 25 def SingleQuote(s): | |
| 26 """Return an shell-escaped version of the string using single quotes. | |
| 27 | |
| 28 Reliably quote a string which may contain unsafe characters (e.g. space, | |
| 29 quote, or other special characters such as '$'). | |
| 30 | |
| 31 The returned value can be used in a shell command line as one token that gets | |
| 32 to be interpreted literally. | |
| 33 | |
| 34 Args: | |
| 35 s: The string to quote. | |
| 36 | |
| 37 Return: | |
| 38 The string quoted using single quotes. | |
| 39 """ | |
| 40 return pipes.quote(s) | |
| 41 | |
| 42 def DoubleQuote(s): | |
| 43 """Return an shell-escaped version of the string using double quotes. | |
| 44 | |
| 45 Reliably quote a string which may contain unsafe characters (e.g. space | |
| 46 or quote characters), while retaining some shell features such as variable | |
| 47 interpolation. | |
| 48 | |
| 49 The returned value can be used in a shell command line as one token that gets | |
| 50 to be further interpreted by the shell. | |
| 51 | |
| 52 The set of characters that retain their special meaning may depend on the | |
| 53 shell implementation. This set usually includes: '$', '`', '\', '!', '*', | |
| 54 and '@'. | |
| 55 | |
| 56 Args: | |
| 57 s: The string to quote. | |
| 58 | |
| 59 Return: | |
| 60 The string quoted using double quotes. | |
| 61 """ | |
| 62 if not s: | |
| 63 return '""' | |
| 64 elif all(c in _SafeShellChars for c in s): | |
| 65 return s | |
| 66 else: | |
| 67 return '"' + s.replace('"', '\\"') + '"' | |
| 68 | |
| 69 | |
| 70 def Popen(args, stdout=None, stderr=None, shell=None, cwd=None, env=None): | |
| 71 return subprocess.Popen( | |
| 72 args=args, cwd=cwd, stdout=stdout, stderr=stderr, | |
| 73 shell=shell, close_fds=True, env=env, | |
| 74 preexec_fn=lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL)) | |
| 75 | |
| 76 | |
| 77 def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None): | |
| 78 pipe = Popen(args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd, | |
| 79 env=env) | |
| 80 pipe.communicate() | |
| 81 return pipe.wait() | |
| 82 | |
| 83 | |
| 84 def RunCmd(args, cwd=None): | |
| 85 """Opens a subprocess to execute a program and returns its return value. | |
| 86 | |
| 87 Args: | |
| 88 args: A string or a sequence of program arguments. The program to execute is | |
| 89 the string or the first item in the args sequence. | |
| 90 cwd: If not None, the subprocess's current directory will be changed to | |
| 91 |cwd| before it's executed. | |
| 92 | |
| 93 Returns: | |
| 94 Return code from the command execution. | |
| 95 """ | |
| 96 logging.info(str(args) + ' ' + (cwd or '')) | |
| 97 return Call(args, cwd=cwd) | |
| 98 | |
| 99 | |
| 100 def GetCmdOutput(args, cwd=None, shell=False): | |
| 101 """Open a subprocess to execute a program and returns its output. | |
| 102 | |
| 103 Args: | |
| 104 args: A string or a sequence of program arguments. The program to execute is | |
| 105 the string or the first item in the args sequence. | |
| 106 cwd: If not None, the subprocess's current directory will be changed to | |
| 107 |cwd| before it's executed. | |
| 108 shell: Whether to execute args as a shell command. | |
| 109 | |
| 110 Returns: | |
| 111 Captures and returns the command's stdout. | |
| 112 Prints the command's stderr to logger (which defaults to stdout). | |
| 113 """ | |
| 114 (_, output) = GetCmdStatusAndOutput(args, cwd, shell) | |
| 115 return output | |
| 116 | |
| 117 | |
| 118 def _ValidateAndLogCommand(args, cwd, shell): | |
| 119 if isinstance(args, basestring): | |
| 120 if not shell: | |
| 121 raise Exception('string args must be run with shell=True') | |
| 122 else: | |
| 123 if shell: | |
| 124 raise Exception('array args must be run with shell=False') | |
| 125 args = ' '.join(SingleQuote(c) for c in args) | |
| 126 if cwd is None: | |
| 127 cwd = '' | |
| 128 else: | |
| 129 cwd = ':' + cwd | |
| 130 logging.info('[host]%s> %s', cwd, args) | |
| 131 return args | |
| 132 | |
| 133 | |
| 134 def GetCmdStatusAndOutput(args, cwd=None, shell=False): | |
| 135 """Executes a subprocess and returns its exit code and output. | |
| 136 | |
| 137 Args: | |
| 138 args: A string or a sequence of program arguments. The program to execute is | |
| 139 the string or the first item in the args sequence. | |
| 140 cwd: If not None, the subprocess's current directory will be changed to | |
| 141 |cwd| before it's executed. | |
| 142 shell: Whether to execute args as a shell command. Must be True if args | |
| 143 is a string and False if args is a sequence. | |
| 144 | |
| 145 Returns: | |
| 146 The 2-tuple (exit code, output). | |
| 147 """ | |
| 148 _ValidateAndLogCommand(args, cwd, shell) | |
| 149 pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, | |
| 150 shell=shell, cwd=cwd) | |
| 151 stdout, stderr = pipe.communicate() | |
| 152 | |
| 153 if stderr: | |
| 154 logging.critical(stderr) | |
| 155 if len(stdout) > 4096: | |
| 156 logging.debug('Truncated output:') | |
| 157 logging.debug(stdout[:4096]) | |
| 158 return (pipe.returncode, stdout) | |
| 159 | |
| 160 | |
| 161 class TimeoutError(Exception): | |
| 162 """Module-specific timeout exception.""" | |
| 163 pass | |
| 164 | |
| 165 | |
| 166 def _IterProcessStdout(process, timeout=None, buffer_size=4096, | |
| 167 poll_interval=1): | |
| 168 assert fcntl, 'fcntl module is required' | |
| 169 try: | |
| 170 # Enable non-blocking reads from the child's stdout. | |
| 171 child_fd = process.stdout.fileno() | |
| 172 fl = fcntl.fcntl(child_fd, fcntl.F_GETFL) | |
| 173 fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) | |
| 174 | |
| 175 end_time = (time.time() + timeout) if timeout else None | |
| 176 while True: | |
| 177 if end_time and time.time() > end_time: | |
| 178 raise TimeoutError | |
| 179 read_fds, _, _ = select.select([child_fd], [], [], poll_interval) | |
| 180 if child_fd in read_fds: | |
| 181 data = os.read(child_fd, buffer_size) | |
| 182 if not data: | |
| 183 break | |
| 184 yield data | |
| 185 if process.poll() is not None: | |
| 186 break | |
| 187 finally: | |
| 188 try: | |
| 189 # Make sure the process doesn't stick around if we fail with an | |
| 190 # exception. | |
| 191 process.kill() | |
| 192 except OSError: | |
| 193 pass | |
| 194 process.wait() | |
| 195 | |
| 196 | |
| 197 def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False, | |
| 198 logfile=None): | |
| 199 """Executes a subprocess with a timeout. | |
| 200 | |
| 201 Args: | |
| 202 args: List of arguments to the program, the program to execute is the first | |
| 203 element. | |
| 204 timeout: the timeout in seconds or None to wait forever. | |
| 205 cwd: If not None, the subprocess's current directory will be changed to | |
| 206 |cwd| before it's executed. | |
| 207 shell: Whether to execute args as a shell command. Must be True if args | |
| 208 is a string and False if args is a sequence. | |
| 209 logfile: Optional file-like object that will receive output from the | |
| 210 command as it is running. | |
| 211 | |
| 212 Returns: | |
| 213 The 2-tuple (exit code, output). | |
| 214 """ | |
| 215 _ValidateAndLogCommand(args, cwd, shell) | |
| 216 output = StringIO.StringIO() | |
| 217 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, | |
| 218 stderr=subprocess.STDOUT) | |
| 219 for data in _IterProcessStdout(process, timeout=timeout): | |
| 220 if logfile: | |
| 221 logfile.write(data) | |
| 222 output.write(data) | |
| 223 return process.returncode, output.getvalue() | |
| 224 | |
| 225 | |
| 226 def IterCmdOutputLines(args, timeout=None, cwd=None, shell=False, | |
| 227 check_status=True): | |
| 228 """Executes a subprocess and continuously yields lines from its output. | |
| 229 | |
| 230 Args: | |
| 231 args: List of arguments to the program, the program to execute is the first | |
| 232 element. | |
| 233 cwd: If not None, the subprocess's current directory will be changed to | |
| 234 |cwd| before it's executed. | |
| 235 shell: Whether to execute args as a shell command. Must be True if args | |
| 236 is a string and False if args is a sequence. | |
| 237 check_status: A boolean indicating whether to check the exit status of the | |
| 238 process after all output has been read. | |
| 239 | |
| 240 Yields: | |
| 241 The output of the subprocess, line by line. | |
| 242 | |
| 243 Raises: | |
| 244 CalledProcessError if check_status is True and the process exited with a | |
| 245 non-zero exit status. | |
| 246 """ | |
| 247 cmd = _ValidateAndLogCommand(args, cwd, shell) | |
| 248 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, | |
| 249 stderr=subprocess.STDOUT) | |
| 250 buffer_output = '' | |
| 251 for data in _IterProcessStdout(process, timeout=timeout): | |
| 252 buffer_output += data | |
| 253 has_incomplete_line = buffer_output[-1] not in '\r\n' | |
| 254 lines = buffer_output.splitlines() | |
| 255 buffer_output = lines.pop() if has_incomplete_line else '' | |
| 256 for line in lines: | |
| 257 yield line | |
| 258 if buffer_output: | |
| 259 yield buffer_output | |
| 260 if check_status and process.returncode: | |
| 261 raise subprocess.CalledProcessError(process.returncode, cmd) | |
| OLD | NEW |