OLD | NEW |
1 # Copyright (c) 2012 The Chromium Authors. All rights reserved. | 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 | 2 # Use of this source code is governed by a BSD-style license that can be |
3 # found in the LICENSE file. | 3 # found in the LICENSE file. |
4 | 4 |
5 """A wrapper for subprocess to make calling shell commands easier.""" | 5 """A wrapper for subprocess to make calling shell commands easier.""" |
6 | 6 |
7 import logging | 7 import logging |
8 import os | 8 import os |
9 import pipes | 9 import pipes |
10 import select | 10 import select |
(...skipping 97 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
108 shell: Whether to execute args as a shell command. | 108 shell: Whether to execute args as a shell command. |
109 | 109 |
110 Returns: | 110 Returns: |
111 Captures and returns the command's stdout. | 111 Captures and returns the command's stdout. |
112 Prints the command's stderr to logger (which defaults to stdout). | 112 Prints the command's stderr to logger (which defaults to stdout). |
113 """ | 113 """ |
114 (_, output) = GetCmdStatusAndOutput(args, cwd, shell) | 114 (_, output) = GetCmdStatusAndOutput(args, cwd, shell) |
115 return output | 115 return output |
116 | 116 |
117 | 117 |
118 def _LogCommand(args, cwd=None): | 118 def _ValidateAndLogCommand(args, cwd, shell): |
119 if not isinstance(args, basestring): | 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') |
120 args = ' '.join(SingleQuote(c) for c in args) | 125 args = ' '.join(SingleQuote(c) for c in args) |
121 if cwd is None: | 126 if cwd is None: |
122 cwd = '' | 127 cwd = '' |
123 else: | 128 else: |
124 cwd = ':' + cwd | 129 cwd = ':' + cwd |
125 logging.info('[host]%s> %s', cwd, args) | 130 logging.info('[host]%s> %s', cwd, args) |
| 131 return args |
126 | 132 |
127 | 133 |
128 def GetCmdStatusAndOutput(args, cwd=None, shell=False): | 134 def GetCmdStatusAndOutput(args, cwd=None, shell=False): |
129 """Executes a subprocess and returns its exit code and output. | 135 """Executes a subprocess and returns its exit code and output. |
130 | 136 |
131 Args: | 137 Args: |
132 args: A string or a sequence of program arguments. The program to execute is | 138 args: A string or a sequence of program arguments. The program to execute is |
133 the string or the first item in the args sequence. | 139 the string or the first item in the args sequence. |
134 cwd: If not None, the subprocess's current directory will be changed to | 140 cwd: If not None, the subprocess's current directory will be changed to |
135 |cwd| before it's executed. | 141 |cwd| before it's executed. |
136 shell: Whether to execute args as a shell command. | 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. |
137 | 144 |
138 Returns: | 145 Returns: |
139 The 2-tuple (exit code, output). | 146 The 2-tuple (exit code, output). |
140 """ | 147 """ |
141 if isinstance(args, basestring): | 148 _ValidateAndLogCommand(args, cwd, shell) |
142 if not shell: | |
143 raise Exception('string args must be run with shell=True') | |
144 elif shell: | |
145 raise Exception('array args must be run with shell=False') | |
146 | |
147 _LogCommand(args, cwd) | |
148 pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, | 149 pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, |
149 shell=shell, cwd=cwd) | 150 shell=shell, cwd=cwd) |
150 stdout, stderr = pipe.communicate() | 151 stdout, stderr = pipe.communicate() |
151 | 152 |
152 if stderr: | 153 if stderr: |
153 logging.critical(stderr) | 154 logging.critical(stderr) |
154 if len(stdout) > 4096: | 155 if len(stdout) > 4096: |
155 logging.debug('Truncated output:') | 156 logging.debug('Truncated output:') |
156 logging.debug(stdout[:4096]) | 157 logging.debug(stdout[:4096]) |
157 return (pipe.returncode, stdout) | 158 return (pipe.returncode, stdout) |
158 | 159 |
159 | 160 |
160 class TimeoutError(Exception): | 161 class TimeoutError(Exception): |
161 """Module-specific timeout exception.""" | 162 """Module-specific timeout exception.""" |
162 pass | 163 pass |
163 | 164 |
164 | 165 |
165 def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False, | 166 def _IterProcessStdout(process, timeout=None, buffer_size=4096, |
166 logfile=None): | 167 poll_interval=1): |
167 """Executes a subprocess with a timeout. | |
168 | |
169 Args: | |
170 args: List of arguments to the program, the program to execute is the first | |
171 element. | |
172 timeout: the timeout in seconds or None to wait forever. | |
173 cwd: If not None, the subprocess's current directory will be changed to | |
174 |cwd| before it's executed. | |
175 shell: Whether to execute args as a shell command. | |
176 logfile: Optional file-like object that will receive output from the | |
177 command as it is running. | |
178 | |
179 Returns: | |
180 The 2-tuple (exit code, output). | |
181 """ | |
182 assert fcntl, 'fcntl module is required' | 168 assert fcntl, 'fcntl module is required' |
183 if isinstance(args, basestring): | |
184 if not shell: | |
185 raise Exception('string args must be run with shell=True') | |
186 elif shell: | |
187 raise Exception('array args must be run with shell=False') | |
188 | |
189 _LogCommand(args, cwd) | |
190 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, | |
191 stderr=subprocess.STDOUT) | |
192 try: | 169 try: |
193 end_time = (time.time() + timeout) if timeout else None | 170 # Enable non-blocking reads from the child's stdout. |
194 poll_interval = 1 | |
195 buffer_size = 4096 | |
196 child_fd = process.stdout.fileno() | 171 child_fd = process.stdout.fileno() |
197 output = StringIO.StringIO() | |
198 | |
199 # Enable non-blocking reads from the child's stdout. | |
200 fl = fcntl.fcntl(child_fd, fcntl.F_GETFL) | 172 fl = fcntl.fcntl(child_fd, fcntl.F_GETFL) |
201 fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) | 173 fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) |
202 | 174 |
| 175 end_time = (time.time() + timeout) if timeout else None |
203 while True: | 176 while True: |
204 if end_time and time.time() > end_time: | 177 if end_time and time.time() > end_time: |
205 raise TimeoutError | 178 raise TimeoutError |
206 read_fds, _, _ = select.select([child_fd], [], [], poll_interval) | 179 read_fds, _, _ = select.select([child_fd], [], [], poll_interval) |
207 if child_fd in read_fds: | 180 if child_fd in read_fds: |
208 data = os.read(child_fd, buffer_size) | 181 data = os.read(child_fd, buffer_size) |
209 if not data: | 182 if not data: |
210 break | 183 break |
211 if logfile: | 184 yield data |
212 logfile.write(data) | |
213 output.write(data) | |
214 if process.poll() is not None: | 185 if process.poll() is not None: |
215 break | 186 break |
216 finally: | 187 finally: |
217 try: | 188 try: |
218 # Make sure the process doesn't stick around if we fail with an | 189 # Make sure the process doesn't stick around if we fail with an |
219 # exception. | 190 # exception. |
220 process.kill() | 191 process.kill() |
221 except OSError: | 192 except OSError: |
222 pass | 193 pass |
223 process.wait() | 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) |
224 return process.returncode, output.getvalue() | 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 |