• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
7import logging
8import os
9import pipes
10import select
11import signal
12import string
13import StringIO
14import subprocess
15import sys
16import time
17
18
19logger = logging.getLogger(__name__)
20
21_SafeShellChars = frozenset(string.ascii_letters + string.digits + '@%_-+=:,./')
22
23
24def SingleQuote(s):
25  """Return an shell-escaped version of the string using single quotes.
26
27  Reliably quote a string which may contain unsafe characters (e.g. space,
28  quote, or other special characters such as '$').
29
30  The returned value can be used in a shell command line as one token that gets
31  to be interpreted literally.
32
33  Args:
34    s: The string to quote.
35
36  Return:
37    The string quoted using single quotes.
38  """
39  return pipes.quote(s)
40
41
42def 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
70def ShrinkToSnippet(cmd_parts, var_name, var_value):
71  """Constructs a shell snippet for a command using a variable to shrink it.
72
73  Takes into account all quoting that needs to happen.
74
75  Args:
76    cmd_parts: A list of command arguments.
77    var_name: The variable that holds var_value.
78    var_value: The string to replace in cmd_parts with $var_name
79
80  Returns:
81    A shell snippet that does not include setting the variable.
82  """
83  def shrink(value):
84    parts = (x and SingleQuote(x) for x in value.split(var_value))
85    with_substitutions = ('"$%s"' % var_name).join(parts)
86    return with_substitutions or "''"
87
88  return ' '.join(shrink(part) for part in cmd_parts)
89
90
91def Popen(args, stdout=None, stderr=None, shell=None, cwd=None, env=None):
92  # preexec_fn isn't supported on windows.
93  if sys.platform == 'win32':
94    close_fds = (stdout is None and stderr is None)
95    preexec_fn = None
96  else:
97    close_fds = True
98    preexec_fn = lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL)
99
100  return subprocess.Popen(
101      args=args, cwd=cwd, stdout=stdout, stderr=stderr,
102      shell=shell, close_fds=close_fds, env=env, preexec_fn=preexec_fn)
103
104
105def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None):
106  pipe = Popen(args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd,
107               env=env)
108  pipe.communicate()
109  return pipe.wait()
110
111
112def RunCmd(args, cwd=None):
113  """Opens a subprocess to execute a program and returns its return value.
114
115  Args:
116    args: A string or a sequence of program arguments. The program to execute is
117      the string or the first item in the args sequence.
118    cwd: If not None, the subprocess's current directory will be changed to
119      |cwd| before it's executed.
120
121  Returns:
122    Return code from the command execution.
123  """
124  logger.info(str(args) + ' ' + (cwd or ''))
125  return Call(args, cwd=cwd)
126
127
128def GetCmdOutput(args, cwd=None, shell=False, env=None):
129  """Open a subprocess to execute a program and returns its output.
130
131  Args:
132    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.
134    cwd: If not None, the subprocess's current directory will be changed to
135      |cwd| before it's executed.
136    shell: Whether to execute args as a shell command.
137    env: If not None, a mapping that defines environment variables for the
138      subprocess.
139
140  Returns:
141    Captures and returns the command's stdout.
142    Prints the command's stderr to logger (which defaults to stdout).
143  """
144  (_, output) = GetCmdStatusAndOutput(args, cwd, shell, env)
145  return output
146
147
148def _ValidateAndLogCommand(args, cwd, shell):
149  if isinstance(args, basestring):
150    if not shell:
151      raise Exception('string args must be run with shell=True')
152  else:
153    if shell:
154      raise Exception('array args must be run with shell=False')
155    args = ' '.join(SingleQuote(str(c)) for c in args)
156  if cwd is None:
157    cwd = ''
158  else:
159    cwd = ':' + cwd
160  logger.debug('[host]%s> %s', cwd, args)
161  return args
162
163
164def GetCmdStatusAndOutput(args, cwd=None, shell=False, env=None):
165  """Executes a subprocess and returns its exit code and output.
166
167  Args:
168    args: A string or a sequence of program arguments. The program to execute is
169      the string or the first item in the args sequence.
170    cwd: If not None, the subprocess's current directory will be changed to
171      |cwd| before it's executed.
172    shell: Whether to execute args as a shell command. Must be True if args
173      is a string and False if args is a sequence.
174    env: If not None, a mapping that defines environment variables for the
175      subprocess.
176
177  Returns:
178    The 2-tuple (exit code, stdout).
179  """
180  status, stdout, stderr = GetCmdStatusOutputAndError(
181      args, cwd=cwd, shell=shell, env=env)
182
183  if stderr:
184    logger.critical('STDERR: %s', stderr)
185  logger.debug('STDOUT: %s%s', stdout[:4096].rstrip(),
186               '<truncated>' if len(stdout) > 4096 else '')
187  return (status, stdout)
188
189
190def StartCmd(args, cwd=None, shell=False, env=None):
191  """Starts a subprocess and returns a handle to the process.
192
193  Args:
194    args: A string or a sequence of program arguments. The program to execute is
195      the string or the first item in the args sequence.
196    cwd: If not None, the subprocess's current directory will be changed to
197      |cwd| before it's executed.
198    shell: Whether to execute args as a shell command. Must be True if args
199      is a string and False if args is a sequence.
200    env: If not None, a mapping that defines environment variables for the
201      subprocess.
202
203  Returns:
204    A process handle from subprocess.Popen.
205  """
206  _ValidateAndLogCommand(args, cwd, shell)
207  return Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
208               shell=shell, cwd=cwd, env=env)
209
210
211def GetCmdStatusOutputAndError(args, cwd=None, shell=False, env=None):
212  """Executes a subprocess and returns its exit code, output, and errors.
213
214  Args:
215    args: A string or a sequence of program arguments. The program to execute is
216      the string or the first item in the args sequence.
217    cwd: If not None, the subprocess's current directory will be changed to
218      |cwd| before it's executed.
219    shell: Whether to execute args as a shell command. Must be True if args
220      is a string and False if args is a sequence.
221    env: If not None, a mapping that defines environment variables for the
222      subprocess.
223
224  Returns:
225    The 3-tuple (exit code, stdout, stderr).
226  """
227  _ValidateAndLogCommand(args, cwd, shell)
228  pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
229               shell=shell, cwd=cwd, env=env)
230  stdout, stderr = pipe.communicate()
231  return (pipe.returncode, stdout, stderr)
232
233
234class TimeoutError(Exception):
235  """Module-specific timeout exception."""
236
237  def __init__(self, output=None):
238    super(TimeoutError, self).__init__()
239    self._output = output
240
241  @property
242  def output(self):
243    return self._output
244
245
246def _IterProcessStdoutFcntl(
247    process, iter_timeout=None, timeout=None, buffer_size=4096,
248    poll_interval=1):
249  """An fcntl-based implementation of _IterProcessStdout."""
250  # pylint: disable=too-many-nested-blocks
251  import fcntl
252  try:
253    # Enable non-blocking reads from the child's stdout.
254    child_fd = process.stdout.fileno()
255    fl = fcntl.fcntl(child_fd, fcntl.F_GETFL)
256    fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
257
258    end_time = (time.time() + timeout) if timeout else None
259    iter_end_time = (time.time() + iter_timeout) if iter_timeout else None
260
261    while True:
262      if end_time and time.time() > end_time:
263        raise TimeoutError()
264      if iter_end_time and time.time() > iter_end_time:
265        yield None
266        iter_end_time = time.time() + iter_timeout
267
268      if iter_end_time:
269        iter_aware_poll_interval = min(
270            poll_interval,
271            max(0, iter_end_time - time.time()))
272      else:
273        iter_aware_poll_interval = poll_interval
274
275      read_fds, _, _ = select.select(
276          [child_fd], [], [], iter_aware_poll_interval)
277      if child_fd in read_fds:
278        data = os.read(child_fd, buffer_size)
279        if not data:
280          break
281        yield data
282
283      if process.poll() is not None:
284        # If process is closed, keep checking for output data (because of timing
285        # issues).
286        while True:
287          read_fds, _, _ = select.select(
288              [child_fd], [], [], iter_aware_poll_interval)
289          if child_fd in read_fds:
290            data = os.read(child_fd, buffer_size)
291            if data:
292              yield data
293              continue
294          break
295        break
296  finally:
297    try:
298      if process.returncode is None:
299        # Make sure the process doesn't stick around if we fail with an
300        # exception.
301        process.kill()
302    except OSError:
303      pass
304    process.wait()
305
306
307def _IterProcessStdoutQueue(
308    process, iter_timeout=None, timeout=None, buffer_size=4096,
309    poll_interval=1):
310  """A Queue.Queue-based implementation of _IterProcessStdout.
311
312  TODO(jbudorick): Evaluate whether this is a suitable replacement for
313  _IterProcessStdoutFcntl on all platforms.
314  """
315  # pylint: disable=unused-argument
316  import Queue
317  import threading
318
319  stdout_queue = Queue.Queue()
320
321  def read_process_stdout():
322    # TODO(jbudorick): Pick an appropriate read size here.
323    while True:
324      try:
325        output_chunk = os.read(process.stdout.fileno(), buffer_size)
326      except IOError:
327        break
328      stdout_queue.put(output_chunk, True)
329      if not output_chunk and process.poll() is not None:
330        break
331
332  reader_thread = threading.Thread(target=read_process_stdout)
333  reader_thread.start()
334
335  end_time = (time.time() + timeout) if timeout else None
336
337  try:
338    while True:
339      if end_time and time.time() > end_time:
340        raise TimeoutError()
341      try:
342        s = stdout_queue.get(True, iter_timeout)
343        if not s:
344          break
345        yield s
346      except Queue.Empty:
347        yield None
348  finally:
349    try:
350      if process.returncode is None:
351        # Make sure the process doesn't stick around if we fail with an
352        # exception.
353        process.kill()
354    except OSError:
355      pass
356    process.wait()
357    reader_thread.join()
358
359
360_IterProcessStdout = (
361    _IterProcessStdoutQueue
362    if sys.platform == 'win32'
363    else _IterProcessStdoutFcntl)
364"""Iterate over a process's stdout.
365
366This is intentionally not public.
367
368Args:
369  process: The process in question.
370  iter_timeout: An optional length of time, in seconds, to wait in
371    between each iteration. If no output is received in the given
372    time, this generator will yield None.
373  timeout: An optional length of time, in seconds, during which
374    the process must finish. If it fails to do so, a TimeoutError
375    will be raised.
376  buffer_size: The maximum number of bytes to read (and thus yield) at once.
377  poll_interval: The length of time to wait in calls to `select.select`.
378    If iter_timeout is set, the remaining length of time in the iteration
379    may take precedence.
380Raises:
381  TimeoutError: if timeout is set and the process does not complete.
382Yields:
383  basestrings of data or None.
384"""
385
386
387def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False,
388                                     logfile=None, env=None):
389  """Executes a subprocess with a timeout.
390
391  Args:
392    args: List of arguments to the program, the program to execute is the first
393      element.
394    timeout: the timeout in seconds or None to wait forever.
395    cwd: If not None, the subprocess's current directory will be changed to
396      |cwd| before it's executed.
397    shell: Whether to execute args as a shell command. Must be True if args
398      is a string and False if args is a sequence.
399    logfile: Optional file-like object that will receive output from the
400      command as it is running.
401    env: If not None, a mapping that defines environment variables for the
402      subprocess.
403
404  Returns:
405    The 2-tuple (exit code, output).
406  Raises:
407    TimeoutError on timeout.
408  """
409  _ValidateAndLogCommand(args, cwd, shell)
410  output = StringIO.StringIO()
411  process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE,
412                  stderr=subprocess.STDOUT, env=env)
413  try:
414    for data in _IterProcessStdout(process, timeout=timeout):
415      if logfile:
416        logfile.write(data)
417      output.write(data)
418  except TimeoutError:
419    raise TimeoutError(output.getvalue())
420
421  str_output = output.getvalue()
422  logger.debug('STDOUT+STDERR: %s%s', str_output[:4096].rstrip(),
423               '<truncated>' if len(str_output) > 4096 else '')
424  return process.returncode, str_output
425
426
427def IterCmdOutputLines(args, iter_timeout=None, timeout=None, cwd=None,
428                       shell=False, env=None, check_status=True):
429  """Executes a subprocess and continuously yields lines from its output.
430
431  Args:
432    args: List of arguments to the program, the program to execute is the first
433      element.
434    iter_timeout: Timeout for each iteration, in seconds.
435    timeout: Timeout for the entire command, in seconds.
436    cwd: If not None, the subprocess's current directory will be changed to
437      |cwd| before it's executed.
438    shell: Whether to execute args as a shell command. Must be True if args
439      is a string and False if args is a sequence.
440    env: If not None, a mapping that defines environment variables for the
441      subprocess.
442    check_status: A boolean indicating whether to check the exit status of the
443      process after all output has been read.
444  Yields:
445    The output of the subprocess, line by line.
446
447  Raises:
448    CalledProcessError if check_status is True and the process exited with a
449      non-zero exit status.
450  """
451  cmd = _ValidateAndLogCommand(args, cwd, shell)
452  process = Popen(args, cwd=cwd, shell=shell, env=env,
453                  stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
454  return _IterCmdOutputLines(
455      process, cmd, iter_timeout=iter_timeout, timeout=timeout,
456      check_status=check_status)
457
458def _IterCmdOutputLines(process, cmd, iter_timeout=None, timeout=None,
459                        check_status=True):
460  buffer_output = ''
461
462  iter_end = None
463  cur_iter_timeout = None
464  if iter_timeout:
465    iter_end = time.time() + iter_timeout
466    cur_iter_timeout = iter_timeout
467
468  for data in _IterProcessStdout(process, iter_timeout=cur_iter_timeout,
469                                 timeout=timeout):
470    if iter_timeout:
471      # Check whether the current iteration has timed out.
472      cur_iter_timeout = iter_end - time.time()
473      if data is None or cur_iter_timeout < 0:
474        yield None
475        iter_end = time.time() + iter_timeout
476        continue
477    else:
478      assert data is not None, (
479          'Iteration received no data despite no iter_timeout being set. '
480          'cmd: %s' % cmd)
481
482    # Construct lines to yield from raw data.
483    buffer_output += data
484    has_incomplete_line = buffer_output[-1] not in '\r\n'
485    lines = buffer_output.splitlines()
486    buffer_output = lines.pop() if has_incomplete_line else ''
487    for line in lines:
488      yield line
489      if iter_timeout:
490        iter_end = time.time() + iter_timeout
491
492  if buffer_output:
493    yield buffer_output
494  if check_status and process.returncode:
495    raise subprocess.CalledProcessError(process.returncode, cmd)
496