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