1# Copyright (c) 2006-2008 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"""Shared process-related utility functions."""
5
6import errno
7import os
8import subprocess
9import sys
10
11class CommandNotFound(Exception): pass
12
13
14TASKKILL = os.path.join(os.environ['WINDIR'], 'system32', 'taskkill.exe')
15TASKKILL_PROCESS_NOT_FOUND_ERR = 128
16# On windows 2000 there is no taskkill.exe, we need to have pskill somewhere
17# in the path.
18PSKILL = 'pskill.exe'
19PSKILL_PROCESS_NOT_FOUND_ERR = -1
20
21def KillAll(executables):
22  """Tries to kill all copies of each process in the processes list.  Returns
23  an error if any running processes couldn't be killed.
24  """
25  result = 0
26  if os.path.exists(TASKKILL):
27    command = [TASKKILL, '/f', '/im']
28    process_not_found_err = TASKKILL_PROCESS_NOT_FOUND_ERR
29  else:
30    command = [PSKILL, '/t']
31    process_not_found_err = PSKILL_PROCESS_NOT_FOUND_ERR
32
33  for name in executables:
34    new_error = RunCommand(command + [name])
35    # Ignore "process not found" error.
36    if new_error != 0 and new_error != process_not_found_err:
37      result = new_error
38  return result
39
40def RunCommandFull(command, verbose=True, collect_output=False,
41                   print_output=True):
42  """Runs the command list.
43
44  Prints the given command (which should be a list of one or more strings).
45  If specified, prints its stderr (and optionally stdout) to stdout,
46  line-buffered, converting line endings to CRLF (see note below).  If
47  specified, collects the output as a list of lines and returns it.  Waits
48  for the command to terminate and returns its status.
49
50  Args:
51    command: the full command to run, as a list of one or more strings
52    verbose: if True, combines all output (stdout and stderr) into stdout.
53             Otherwise, prints only the command's stderr to stdout.
54    collect_output: if True, collects the output of the command as a list of
55                    lines and returns it
56    print_output: if True, prints the output of the command
57
58  Returns:
59    A tuple consisting of the process's exit status and output.  If
60    collect_output is False, the output will be [].
61
62  Raises:
63    CommandNotFound if the command executable could not be found.
64  """
65  print '\n' + subprocess.list2cmdline(command).replace('\\', '/') + '\n', ###
66
67  if verbose:
68    out = subprocess.PIPE
69    err = subprocess.STDOUT
70  else:
71    out = file(os.devnull, 'w')
72    err = subprocess.PIPE
73  try:
74    proc = subprocess.Popen(command, stdout=out, stderr=err, bufsize=1)
75  except OSError, e:
76    if e.errno == errno.ENOENT:
77      raise CommandNotFound('Unable to find "%s"' % command[0])
78    raise
79
80  output = []
81
82  if verbose:
83    read_from = proc.stdout
84  else:
85    read_from = proc.stderr
86  line = read_from.readline()
87  while line:
88    line = line.rstrip()
89
90    if collect_output:
91      output.append(line)
92
93    if print_output:
94      # Windows Python converts \n to \r\n automatically whenever it
95      # encounters it written to a text file (including stdout).  The only
96      # way around it is to write to a binary file, which isn't feasible for
97      # stdout. So we end up with \r\n here even though we explicitly write
98      # \n.  (We could write \r instead, which doesn't get converted to \r\n,
99      # but that's probably more troublesome for people trying to read the
100      # files.)
101      print line + '\n',
102
103      # Python on windows writes the buffer only when it reaches 4k. This is
104      # not fast enough for all purposes.
105      sys.stdout.flush()
106    line = read_from.readline()
107
108  # Make sure the process terminates.
109  proc.wait()
110
111  if not verbose:
112    out.close()
113  return (proc.returncode, output)
114
115def RunCommand(command, verbose=True):
116  """Runs the command list, printing its output and returning its exit status.
117
118  Prints the given command (which should be a list of one or more strings),
119  then runs it and prints its stderr (and optionally stdout) to stdout,
120  line-buffered, converting line endings to CRLF.  Waits for the command to
121  terminate and returns its status.
122
123  Args:
124    command: the full command to run, as a list of one or more strings
125    verbose: if True, combines all output (stdout and stderr) into stdout.
126             Otherwise, prints only the command's stderr to stdout.
127
128  Returns:
129    The process's exit status.
130
131  Raises:
132    CommandNotFound if the command executable could not be found.
133  """
134  return RunCommandFull(command, verbose)[0]
135
136def RunCommandsInParallel(commands, verbose=True, collect_output=False,
137                          print_output=True):
138  """Runs a list of commands in parallel, waits for all commands to terminate
139  and returns their status. If specified, the ouput of commands can be
140  returned and/or printed.
141
142  Args:
143    commands: the list of commands to run, each as a list of one or more
144              strings.
145    verbose: if True, combines stdout and stderr into stdout.
146             Otherwise, prints only the command's stderr to stdout.
147    collect_output: if True, collects the output of the each command as a list
148                    of lines and returns it.
149    print_output: if True, prints the output of each command.
150
151  Returns:
152    A list of tuples consisting of each command's exit status and output.  If
153    collect_output is False, the output will be [].
154
155  Raises:
156    CommandNotFound if any of the command executables could not be found.
157  """
158
159  command_num = len(commands)
160  outputs = [[] for i in xrange(command_num)]
161  procs = [None for i in xrange(command_num)]
162  eofs = [False for i in xrange(command_num)]
163
164  for command in commands:
165    print '\n' + subprocess.list2cmdline(command).replace('\\', '/') + '\n',
166
167  if verbose:
168    out = subprocess.PIPE
169    err = subprocess.STDOUT
170  else:
171    out = file(os.devnull, 'w')
172    err = subprocess.PIPE
173
174  for i in xrange(command_num):
175    try:
176      command = commands[i]
177      procs[i] = subprocess.Popen(command, stdout=out, stderr=err, bufsize=1)
178    except OSError, e:
179      if e.errno == errno.ENOENT:
180        raise CommandNotFound('Unable to find "%s"' % command[0])
181      raise
182      # We could consider terminating the processes already started.
183      # But Popen.kill() is only available in version 2.6.
184      # For now the clean up is done by KillAll.
185
186  while True:
187    eof_all = True
188    for i in xrange(command_num):
189      if eofs[i]:
190        continue
191      if verbose:
192        read_from = procs[i].stdout
193      else:
194        read_from = procs[i].stderr
195      line = read_from.readline()
196      if line:
197        eof_all = False
198        line = line.rstrip()
199        outputs[i].append(line)
200        if print_output:
201          # Windows Python converts \n to \r\n automatically whenever it
202          # encounters it written to a text file (including stdout).  The only
203          # way around it is to write to a binary file, which isn't feasible
204          # for stdout. So we end up with \r\n here even though we explicitly
205          # write \n.  (We could write \r instead, which doesn't get converted
206          # to \r\n, but that's probably more troublesome for people trying to
207          # read the files.)
208          print line + '\n',
209      else:
210        eofs[i] = True
211    if eof_all:
212      break
213
214  # Make sure the process terminates.
215  for i in xrange(command_num):
216    procs[i].wait()
217
218  if not verbose:
219    out.close()
220
221  return [(procs[i].returncode, outputs[i]) for i in xrange(command_num)]
222