xref: /aosp_15_r20/external/autotest/utils/frozen_chromite/lib/operation.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1*9c5db199SXin Li# -*- coding: utf-8 -*-
2*9c5db199SXin Li# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
3*9c5db199SXin Li# Use of this source code is governed by a BSD-style license that can be
4*9c5db199SXin Li# found in the LICENSE file.
5*9c5db199SXin Li
6*9c5db199SXin Li"""Operation, including output and progress display
7*9c5db199SXin Li
8*9c5db199SXin LiThis module implements the concept of an operation, which has regular progress
9*9c5db199SXin Liupdates, verbose text display and perhaps some errors.
10*9c5db199SXin Li"""
11*9c5db199SXin Li
12*9c5db199SXin Lifrom __future__ import division
13*9c5db199SXin Lifrom __future__ import print_function
14*9c5db199SXin Li
15*9c5db199SXin Liimport collections
16*9c5db199SXin Liimport contextlib
17*9c5db199SXin Liimport fcntl
18*9c5db199SXin Liimport multiprocessing
19*9c5db199SXin Liimport os
20*9c5db199SXin Liimport pty
21*9c5db199SXin Liimport re
22*9c5db199SXin Liimport struct
23*9c5db199SXin Liimport sys
24*9c5db199SXin Liimport termios
25*9c5db199SXin Li
26*9c5db199SXin Lifrom six.moves import queue as Queue
27*9c5db199SXin Li
28*9c5db199SXin Lifrom autotest_lib.utils.frozen_chromite.lib import cros_logging as logging
29*9c5db199SXin Lifrom autotest_lib.utils.frozen_chromite.lib import osutils
30*9c5db199SXin Lifrom autotest_lib.utils.frozen_chromite.lib import parallel
31*9c5db199SXin Lifrom autotest_lib.utils.frozen_chromite.lib.terminal import Color
32*9c5db199SXin Lifrom autotest_lib.utils.frozen_chromite.utils import outcap
33*9c5db199SXin Li
34*9c5db199SXin Li
35*9c5db199SXin Li# Define filenames for captured stdout and stderr.
36*9c5db199SXin LiSTDOUT_FILE = 'stdout'
37*9c5db199SXin LiSTDERR_FILE = 'stderr'
38*9c5db199SXin Li
39*9c5db199SXin Li_TerminalSize = collections.namedtuple('_TerminalSize', ('lines', 'columns'))
40*9c5db199SXin Li
41*9c5db199SXin Li
42*9c5db199SXin Liclass _BackgroundTaskComplete(object):
43*9c5db199SXin Li  """Sentinal object to indicate that the background task is complete."""
44*9c5db199SXin Li
45*9c5db199SXin Li
46*9c5db199SXin Liclass ProgressBarOperation(object):
47*9c5db199SXin Li  """Wrapper around long running functions to show progress.
48*9c5db199SXin Li
49*9c5db199SXin Li  This class is intended to capture the output of a long running fuction, parse
50*9c5db199SXin Li  the output, and display a progress bar.
51*9c5db199SXin Li
52*9c5db199SXin Li  To display a progress bar for a function foo with argument foo_args, this is
53*9c5db199SXin Li  the usage case:
54*9c5db199SXin Li    1) Create a class that inherits from ProgressBarOperation (e.g.
55*9c5db199SXin Li    FooTypeOperation. In this class, override the ParseOutput method to parse
56*9c5db199SXin Li    the output of foo.
57*9c5db199SXin Li    2) op = operation.FooTypeOperation()
58*9c5db199SXin Li       op.Run(foo, foo_args)
59*9c5db199SXin Li  """
60*9c5db199SXin Li
61*9c5db199SXin Li  # Subtract 10 characters from the width of the terminal because these are used
62*9c5db199SXin Li  # to display the percentage as well as other spaces.
63*9c5db199SXin Li  _PROGRESS_BAR_BORDER_SIZE = 10
64*9c5db199SXin Li
65*9c5db199SXin Li  # By default, update the progress bar every 100 ms.
66*9c5db199SXin Li  _PROGRESS_BAR_UPDATE_INTERVAL = 0.1
67*9c5db199SXin Li
68*9c5db199SXin Li  def __init__(self):
69*9c5db199SXin Li    self._queue = multiprocessing.Queue()
70*9c5db199SXin Li    self._stderr = None
71*9c5db199SXin Li    self._stdout = None
72*9c5db199SXin Li    self._stdout_path = None
73*9c5db199SXin Li    self._stderr_path = None
74*9c5db199SXin Li    self._progress_bar_displayed = False
75*9c5db199SXin Li    self._isatty = os.isatty(sys.stdout.fileno())
76*9c5db199SXin Li
77*9c5db199SXin Li  def _GetTerminalSize(self, fd=pty.STDOUT_FILENO):
78*9c5db199SXin Li    """Return a terminal size object for |fd|.
79*9c5db199SXin Li
80*9c5db199SXin Li    Note: Replace with os.terminal_size() in python3.3.
81*9c5db199SXin Li    """
82*9c5db199SXin Li    winsize = struct.pack('HHHH', 0, 0, 0, 0)
83*9c5db199SXin Li    data = fcntl.ioctl(fd, termios.TIOCGWINSZ, winsize)
84*9c5db199SXin Li    winsize = struct.unpack('HHHH', data)
85*9c5db199SXin Li    return _TerminalSize(int(winsize[0]), int(winsize[1]))
86*9c5db199SXin Li
87*9c5db199SXin Li  def ProgressBar(self, progress):
88*9c5db199SXin Li    """This method creates and displays a progress bar.
89*9c5db199SXin Li
90*9c5db199SXin Li    If not in a terminal, we do not display a progress bar.
91*9c5db199SXin Li
92*9c5db199SXin Li    Args:
93*9c5db199SXin Li      progress: a float between 0 and 1 that represents the fraction of the
94*9c5db199SXin Li        current progress.
95*9c5db199SXin Li    """
96*9c5db199SXin Li    if not self._isatty:
97*9c5db199SXin Li      return
98*9c5db199SXin Li    self._progress_bar_displayed = True
99*9c5db199SXin Li    progress = max(0.0, min(1.0, progress))
100*9c5db199SXin Li    width = max(1, self._GetTerminalSize().columns -
101*9c5db199SXin Li                self._PROGRESS_BAR_BORDER_SIZE)
102*9c5db199SXin Li    block = int(width * progress)
103*9c5db199SXin Li    shaded = '#' * block
104*9c5db199SXin Li    unshaded = '-' * (width - block)
105*9c5db199SXin Li    text = '\r [%s%s] %d%%' % (shaded, unshaded, progress * 100)
106*9c5db199SXin Li    sys.stdout.write(text)
107*9c5db199SXin Li    sys.stdout.flush()
108*9c5db199SXin Li
109*9c5db199SXin Li  def OpenStdoutStderr(self):
110*9c5db199SXin Li    """Open the stdout and stderr streams."""
111*9c5db199SXin Li    if self._stdout is None and self._stderr is None:
112*9c5db199SXin Li      self._stdout = open(self._stdout_path, 'r')
113*9c5db199SXin Li      self._stderr = open(self._stderr_path, 'r')
114*9c5db199SXin Li
115*9c5db199SXin Li  def Cleanup(self):
116*9c5db199SXin Li    """Method to cleanup progress bar.
117*9c5db199SXin Li
118*9c5db199SXin Li    If progress bar has been printed, then we make sure it displays 100% before
119*9c5db199SXin Li    exiting.
120*9c5db199SXin Li    """
121*9c5db199SXin Li    if self._progress_bar_displayed:
122*9c5db199SXin Li      self.ProgressBar(1)
123*9c5db199SXin Li      sys.stdout.write('\n')
124*9c5db199SXin Li      sys.stdout.flush()
125*9c5db199SXin Li
126*9c5db199SXin Li  def ParseOutput(self, output=None):
127*9c5db199SXin Li    """Method to parse output and update progress bar.
128*9c5db199SXin Li
129*9c5db199SXin Li    This method should be overridden to read and parse the lines in _stdout and
130*9c5db199SXin Li    _stderr.
131*9c5db199SXin Li
132*9c5db199SXin Li    One example use of this method could be to detect 'foo' in stdout and
133*9c5db199SXin Li    increment the progress bar every time foo is seen.
134*9c5db199SXin Li
135*9c5db199SXin Li    def ParseOutput(self):
136*9c5db199SXin Li      stdout = self._stdout.read()
137*9c5db199SXin Li      if 'foo' in stdout:
138*9c5db199SXin Li        # Increment progress bar.
139*9c5db199SXin Li
140*9c5db199SXin Li    Args:
141*9c5db199SXin Li      output: Pass in output to parse instead of reading from self._stdout and
142*9c5db199SXin Li        self._stderr.
143*9c5db199SXin Li    """
144*9c5db199SXin Li    raise NotImplementedError('Subclass must override this method.')
145*9c5db199SXin Li
146*9c5db199SXin Li  # TODO(ralphnathan): Deprecate this function and use parallel._BackgroundTask
147*9c5db199SXin Li  # instead (brbug.com/863)
148*9c5db199SXin Li  def WaitUntilComplete(self, update_period):
149*9c5db199SXin Li    """Return True if running background task has completed."""
150*9c5db199SXin Li    try:
151*9c5db199SXin Li      x = self._queue.get(timeout=update_period)
152*9c5db199SXin Li      if isinstance(x, _BackgroundTaskComplete):
153*9c5db199SXin Li        return True
154*9c5db199SXin Li    except Queue.Empty:
155*9c5db199SXin Li      return False
156*9c5db199SXin Li
157*9c5db199SXin Li  def CaptureOutputInBackground(self, func, *args, **kwargs):
158*9c5db199SXin Li    """Launch func in background and capture its output.
159*9c5db199SXin Li
160*9c5db199SXin Li    Args:
161*9c5db199SXin Li      func: Function to execute in the background and whose output is to be
162*9c5db199SXin Li        captured.
163*9c5db199SXin Li      log_level: Logging level to run the func at. By default, it runs at log
164*9c5db199SXin Li        level info.
165*9c5db199SXin Li    """
166*9c5db199SXin Li    log_level = kwargs.pop('log_level', logging.INFO)
167*9c5db199SXin Li    restore_log_level = logging.getLogger().getEffectiveLevel()
168*9c5db199SXin Li    logging.getLogger().setLevel(log_level)
169*9c5db199SXin Li    try:
170*9c5db199SXin Li      with outcap.OutputCapturer(
171*9c5db199SXin Li          stdout_path=self._stdout_path, stderr_path=self._stderr_path,
172*9c5db199SXin Li          quiet_fail=False):
173*9c5db199SXin Li        func(*args, **kwargs)
174*9c5db199SXin Li    finally:
175*9c5db199SXin Li      self._queue.put(_BackgroundTaskComplete())
176*9c5db199SXin Li      logging.getLogger().setLevel(restore_log_level)
177*9c5db199SXin Li
178*9c5db199SXin Li  # TODO (ralphnathan): Store PID of spawned process.
179*9c5db199SXin Li  def Run(self, func, *args, **kwargs):
180*9c5db199SXin Li    """Run func, parse its output, and update the progress bar.
181*9c5db199SXin Li
182*9c5db199SXin Li    Args:
183*9c5db199SXin Li      func: Function to execute in the background and whose output is to be
184*9c5db199SXin Li        captured.
185*9c5db199SXin Li      update_period: Optional argument to specify the period that output should
186*9c5db199SXin Li        be read.
187*9c5db199SXin Li      log_level: Logging level to run the func at. By default, it runs at log
188*9c5db199SXin Li        level info.
189*9c5db199SXin Li    """
190*9c5db199SXin Li    update_period = kwargs.pop('update_period',
191*9c5db199SXin Li                               self._PROGRESS_BAR_UPDATE_INTERVAL)
192*9c5db199SXin Li
193*9c5db199SXin Li    # If we are not running in a terminal device, do not display the progress
194*9c5db199SXin Li    # bar.
195*9c5db199SXin Li    if not self._isatty:
196*9c5db199SXin Li      log_level = kwargs.pop('log_level', logging.INFO)
197*9c5db199SXin Li      restore_log_level = logging.getLogger().getEffectiveLevel()
198*9c5db199SXin Li      logging.getLogger().setLevel(log_level)
199*9c5db199SXin Li      try:
200*9c5db199SXin Li        func(*args, **kwargs)
201*9c5db199SXin Li      finally:
202*9c5db199SXin Li        logging.getLogger().setLevel(restore_log_level)
203*9c5db199SXin Li      return
204*9c5db199SXin Li
205*9c5db199SXin Li    with osutils.TempDir() as tempdir:
206*9c5db199SXin Li      self._stdout_path = os.path.join(tempdir, STDOUT_FILE)
207*9c5db199SXin Li      self._stderr_path = os.path.join(tempdir, STDERR_FILE)
208*9c5db199SXin Li      osutils.Touch(self._stdout_path)
209*9c5db199SXin Li      osutils.Touch(self._stderr_path)
210*9c5db199SXin Li      try:
211*9c5db199SXin Li        with parallel.BackgroundTaskRunner(
212*9c5db199SXin Li            self.CaptureOutputInBackground, func, *args, **kwargs) as queue:
213*9c5db199SXin Li          queue.put([])
214*9c5db199SXin Li          self.OpenStdoutStderr()
215*9c5db199SXin Li          while True:
216*9c5db199SXin Li            self.ParseOutput()
217*9c5db199SXin Li            if self.WaitUntilComplete(update_period):
218*9c5db199SXin Li              break
219*9c5db199SXin Li        # Before we exit, parse the output again to update progress bar.
220*9c5db199SXin Li        self.ParseOutput()
221*9c5db199SXin Li        # Final sanity check to update the progress bar to 100% if it was used
222*9c5db199SXin Li        # by ParseOutput
223*9c5db199SXin Li        self.Cleanup()
224*9c5db199SXin Li      except:
225*9c5db199SXin Li        # Add a blank line before the logging message so the message isn't
226*9c5db199SXin Li        # touching the progress bar.
227*9c5db199SXin Li        sys.stdout.write('\n')
228*9c5db199SXin Li        logging.error('Oops. Something went wrong.')
229*9c5db199SXin Li        # Raise the exception so it can be caught again.
230*9c5db199SXin Li        raise
231*9c5db199SXin Li
232*9c5db199SXin Li
233*9c5db199SXin Liclass ParallelEmergeOperation(ProgressBarOperation):
234*9c5db199SXin Li  """ProgressBarOperation specific for scripts/parallel_emerge.py."""
235*9c5db199SXin Li
236*9c5db199SXin Li  def __init__(self):
237*9c5db199SXin Li    super(ParallelEmergeOperation, self).__init__()
238*9c5db199SXin Li    self._total = None
239*9c5db199SXin Li    self._completed = 0
240*9c5db199SXin Li    self._printed_no_packages = False
241*9c5db199SXin Li    self._events = ['Fetched ', 'Completed ']
242*9c5db199SXin Li    self._msg = None
243*9c5db199SXin Li
244*9c5db199SXin Li  def _GetTotal(self, output):
245*9c5db199SXin Li    """Get total packages by looking for Total: digits packages."""
246*9c5db199SXin Li    match = re.search(r'Total: (\d+) packages', output)
247*9c5db199SXin Li    return int(match.group(1)) if match else None
248*9c5db199SXin Li
249*9c5db199SXin Li  def SetProgressBarMessage(self, msg):
250*9c5db199SXin Li    """Message to be shown before the progress bar is displayed with 0%.
251*9c5db199SXin Li
252*9c5db199SXin Li       The message is not displayed if the progress bar is not going to be
253*9c5db199SXin Li       displayed.
254*9c5db199SXin Li    """
255*9c5db199SXin Li    self._msg = msg
256*9c5db199SXin Li
257*9c5db199SXin Li  def ParseOutput(self, output=None):
258*9c5db199SXin Li    """Parse the output of emerge to determine how to update progress bar.
259*9c5db199SXin Li
260*9c5db199SXin Li    1) Figure out how many packages exist. If the total number of packages to be
261*9c5db199SXin Li    built is zero, then we do not display the progress bar.
262*9c5db199SXin Li    2) Whenever a package is downloaded or built, 'Fetched' and 'Completed' are
263*9c5db199SXin Li    printed respectively. By counting counting 'Fetched's and 'Completed's, we
264*9c5db199SXin Li    can determine how much to update the progress bar by.
265*9c5db199SXin Li
266*9c5db199SXin Li    Args:
267*9c5db199SXin Li      output: Pass in output to parse instead of reading from self._stdout and
268*9c5db199SXin Li        self._stderr.
269*9c5db199SXin Li
270*9c5db199SXin Li    Returns:
271*9c5db199SXin Li      A fraction between 0 and 1 indicating the level of the progress bar. If
272*9c5db199SXin Li      the progress bar isn't displayed, then the return value is -1.
273*9c5db199SXin Li    """
274*9c5db199SXin Li    if output is None:
275*9c5db199SXin Li      stdout = self._stdout.read()
276*9c5db199SXin Li      stderr = self._stderr.read()
277*9c5db199SXin Li      output = stdout + stderr
278*9c5db199SXin Li
279*9c5db199SXin Li    if self._total is None:
280*9c5db199SXin Li      temp = self._GetTotal(output)
281*9c5db199SXin Li      if temp is not None:
282*9c5db199SXin Li        self._total = temp * len(self._events)
283*9c5db199SXin Li        if self._msg is not None:
284*9c5db199SXin Li          logging.notice(self._msg)
285*9c5db199SXin Li
286*9c5db199SXin Li    for event in self._events:
287*9c5db199SXin Li      self._completed += output.count(event)
288*9c5db199SXin Li
289*9c5db199SXin Li    if not self._printed_no_packages and self._total == 0:
290*9c5db199SXin Li      logging.notice('No packages to build.')
291*9c5db199SXin Li      self._printed_no_packages = True
292*9c5db199SXin Li
293*9c5db199SXin Li    if self._total:
294*9c5db199SXin Li      progress = self._completed / self._total
295*9c5db199SXin Li      self.ProgressBar(progress)
296*9c5db199SXin Li      return progress
297*9c5db199SXin Li    else:
298*9c5db199SXin Li      return -1
299*9c5db199SXin Li
300*9c5db199SXin Li
301*9c5db199SXin Li# TODO(sjg): When !isatty(), keep stdout and stderr separate so they can be
302*9c5db199SXin Li# redirected separately
303*9c5db199SXin Li# TODO(sjg): Add proper docs to this fileno
304*9c5db199SXin Li# TODO(sjg): Handle stdin wait in quite mode, rather than silently stalling
305*9c5db199SXin Li
306*9c5db199SXin Liclass Operation(object):
307*9c5db199SXin Li  """Class which controls stdio and progress of an operation in progress.
308*9c5db199SXin Li
309*9c5db199SXin Li  This class is created to handle stdio for a running subprocess. It filters
310*9c5db199SXin Li  it looking for errors and progress information. Optionally it can output the
311*9c5db199SXin Li  stderr and stdout to the terminal, but it is normally supressed.
312*9c5db199SXin Li
313*9c5db199SXin Li  Progress information is garnered from the subprocess output based on
314*9c5db199SXin Li  knowledge of the legacy scripts, but at some point will move over to using
315*9c5db199SXin Li  real progress information reported through new python methods which will
316*9c5db199SXin Li  replace the scripts.
317*9c5db199SXin Li
318*9c5db199SXin Li  Each operation has a name, and this class handles displaying this name
319*9c5db199SXin Li  as it reports progress.
320*9c5db199SXin Li
321*9c5db199SXin Li  Operation Objects
322*9c5db199SXin Li  =================
323*9c5db199SXin Li
324*9c5db199SXin Li  verbose: True / False
325*9c5db199SXin Li    In verbose mode all output from subprocesses is displayed, otherwise
326*9c5db199SXin Li    this output is normally supressed, unless we think it indicates an error.
327*9c5db199SXin Li
328*9c5db199SXin Li  progress: True / False
329*9c5db199SXin Li    The output from subprocesses can be analysed in a very basic manner to
330*9c5db199SXin Li    try to present progress information to the user.
331*9c5db199SXin Li
332*9c5db199SXin Li  explicit_verbose: True / False
333*9c5db199SXin Li    False if we are not just using default verbosity. In that case we allow
334*9c5db199SXin Li    verbosity to be enabled on request, since the user has not explicitly
335*9c5db199SXin Li    disabled it. This is used by commands that the user issues with the
336*9c5db199SXin Li    expectation that output would ordinarily be visible.
337*9c5db199SXin Li  """
338*9c5db199SXin Li
339*9c5db199SXin Li  def __init__(self, name, color=None):
340*9c5db199SXin Li    """Create a new operation.
341*9c5db199SXin Li
342*9c5db199SXin Li    Args:
343*9c5db199SXin Li      name: Operation name in a form to be displayed for the user.
344*9c5db199SXin Li      color: Determines policy for sending color to stdout; see terminal.Color
345*9c5db199SXin Li        for details on interpretation on the value.
346*9c5db199SXin Li    """
347*9c5db199SXin Li    self._name = name   # Operation name.
348*9c5db199SXin Li    self.verbose = False   # True to echo subprocess output.
349*9c5db199SXin Li    self.progress = True   # True to report progress of the operation
350*9c5db199SXin Li    self._column = 0    # Current output column (always 0 unless verbose).
351*9c5db199SXin Li    self._update_len = 0    # Length of last progress update message.
352*9c5db199SXin Li    self._line = ''   # text of current line, so far
353*9c5db199SXin Li    self.explicit_verbose = False
354*9c5db199SXin Li
355*9c5db199SXin Li    self._color = Color(enabled=color)
356*9c5db199SXin Li
357*9c5db199SXin Li    # -1 = no newline pending
358*9c5db199SXin Li    #  n = newline pending, and line length of last line was n
359*9c5db199SXin Li    self._pending_nl = -1
360*9c5db199SXin Li
361*9c5db199SXin Li    # the type of the last stream to emit data on the current lines
362*9c5db199SXin Li    # can be sys.stdout, sys.stderr (both from the subprocess), or None
363*9c5db199SXin Li    # for our own mesages
364*9c5db199SXin Li    self._cur_stream = None
365*9c5db199SXin Li
366*9c5db199SXin Li    self._error_count = 0   # number of error lines we have reported
367*9c5db199SXin Li
368*9c5db199SXin Li  def __del__(self):
369*9c5db199SXin Li    """Object is about to be destroyed, so finish out output cleanly."""
370*9c5db199SXin Li    self.FinishOutput()
371*9c5db199SXin Li
372*9c5db199SXin Li  def FinishOutput(self):
373*9c5db199SXin Li    """Finish off any pending output.
374*9c5db199SXin Li
375*9c5db199SXin Li    This finishes any output line currently in progress and resets the color
376*9c5db199SXin Li    back to normal.
377*9c5db199SXin Li    """
378*9c5db199SXin Li    self._FinishLine(self.verbose, final=True)
379*9c5db199SXin Li    if self._column and self.verbose:
380*9c5db199SXin Li      print(self._color.Stop())
381*9c5db199SXin Li      self._column = 0
382*9c5db199SXin Li
383*9c5db199SXin Li  def WereErrorsDetected(self):
384*9c5db199SXin Li    """Returns whether any errors have been detected.
385*9c5db199SXin Li
386*9c5db199SXin Li    Returns:
387*9c5db199SXin Li      True if any errors have been detected in subprocess output so far.
388*9c5db199SXin Li      False otherwise
389*9c5db199SXin Li    """
390*9c5db199SXin Li    return self._error_count > 0
391*9c5db199SXin Li
392*9c5db199SXin Li  def SetName(self, name):
393*9c5db199SXin Li    """Set the name of the operation as displayed to the user.
394*9c5db199SXin Li
395*9c5db199SXin Li    Args:
396*9c5db199SXin Li      name: Operation name.
397*9c5db199SXin Li    """
398*9c5db199SXin Li    self._name = name
399*9c5db199SXin Li
400*9c5db199SXin Li  def _FilterOutputForErrors(self, line, print_error):
401*9c5db199SXin Li    """Filter a line of output to look for and display errors.
402*9c5db199SXin Li
403*9c5db199SXin Li    This uses a few regular expression searches to spot common error reports
404*9c5db199SXin Li    from subprocesses. A count of these is kept so we know how many occurred.
405*9c5db199SXin Li    Optionally they are displayed in red on the terminal.
406*9c5db199SXin Li
407*9c5db199SXin Li    Args:
408*9c5db199SXin Li      line: the output line to filter, as a string.
409*9c5db199SXin Li      print_error: True to print the error, False to just record it.
410*9c5db199SXin Li    """
411*9c5db199SXin Li    bad_things = ['Cannot GET', 'ERROR', '!!!', 'FAILED']
412*9c5db199SXin Li    for bad_thing in bad_things:
413*9c5db199SXin Li      if re.search(bad_thing, line, flags=re.IGNORECASE):
414*9c5db199SXin Li        self._error_count += 1
415*9c5db199SXin Li        if print_error:
416*9c5db199SXin Li          print(self._color.Color(self._color.RED, line))
417*9c5db199SXin Li          break
418*9c5db199SXin Li
419*9c5db199SXin Li  def _FilterOutputForProgress(self, line):
420*9c5db199SXin Li    """Filter a line of output to look for and dispay progress information.
421*9c5db199SXin Li
422*9c5db199SXin Li    This uses a simple regular expression search to spot progress information
423*9c5db199SXin Li    coming from subprocesses. This is sent to the _Progress() method.
424*9c5db199SXin Li
425*9c5db199SXin Li    Args:
426*9c5db199SXin Li      line: the output line to filter, as a string.
427*9c5db199SXin Li    """
428*9c5db199SXin Li    match = re.match(r'Pending (\d+).*Total (\d+)', line)
429*9c5db199SXin Li    if match:
430*9c5db199SXin Li      pending = int(match.group(1))
431*9c5db199SXin Li      total = int(match.group(2))
432*9c5db199SXin Li      self._Progress(total - pending, total)
433*9c5db199SXin Li
434*9c5db199SXin Li  def _Progress(self, upto, total):
435*9c5db199SXin Li    """Record and optionally display progress information.
436*9c5db199SXin Li
437*9c5db199SXin Li    Args:
438*9c5db199SXin Li      upto: which step we are up to in the operation (integer, from 0).
439*9c5db199SXin Li      total: total number of steps in operation,
440*9c5db199SXin Li    """
441*9c5db199SXin Li    if total > 0:
442*9c5db199SXin Li      update_str = '%s...%d%% (%d of %d)' % (self._name,
443*9c5db199SXin Li                                             upto * 100 // total, upto, total)
444*9c5db199SXin Li      if self.progress:
445*9c5db199SXin Li        # Finish the current line, print progress, and remember its length.
446*9c5db199SXin Li        self._FinishLine(self.verbose)
447*9c5db199SXin Li
448*9c5db199SXin Li        # Sometimes the progress string shrinks and in this case we need to
449*9c5db199SXin Li        # blank out the characters at the end of the line that will not be
450*9c5db199SXin Li        # overwritten by the new line
451*9c5db199SXin Li        pad = max(self._update_len - len(update_str), 0)
452*9c5db199SXin Li        sys.stdout.write(update_str + (' ' * pad) + '\r')
453*9c5db199SXin Li        self._update_len = len(update_str)
454*9c5db199SXin Li
455*9c5db199SXin Li  def _FinishLine(self, display, final=False):
456*9c5db199SXin Li    """Finish off the current line and prepare to start a new one.
457*9c5db199SXin Li
458*9c5db199SXin Li    If a new line is pending from the previous line, then this will be output,
459*9c5db199SXin Li    along with a color reset if needed.
460*9c5db199SXin Li
461*9c5db199SXin Li    We also handle removing progress messages from the output. This is done
462*9c5db199SXin Li    using a carriage return character, following by spaces.
463*9c5db199SXin Li
464*9c5db199SXin Li    Args:
465*9c5db199SXin Li      display: True to display output, False to suppress it
466*9c5db199SXin Li      final: True if this is the final output before we exit, in which case
467*9c5db199SXin Li          we must clean up any remaining progress message by overwriting
468*9c5db199SXin Li          it with spaces, then carriage return
469*9c5db199SXin Li    """
470*9c5db199SXin Li    if display:
471*9c5db199SXin Li      if self._pending_nl != -1:
472*9c5db199SXin Li        # If out last output line was shorter than the progress info
473*9c5db199SXin Li        # add spaces.
474*9c5db199SXin Li        if self._pending_nl < self._update_len:
475*9c5db199SXin Li          print(' ' * (self._update_len - self._pending_nl), end='')
476*9c5db199SXin Li
477*9c5db199SXin Li        # Output the newline, and reset our counter.
478*9c5db199SXin Li        sys.stdout.write(self._color.Stop())
479*9c5db199SXin Li        print()
480*9c5db199SXin Li
481*9c5db199SXin Li    # If this is the last thing that this operation will print, we need to
482*9c5db199SXin Li    # close things off. So if there is some text on the current line but not
483*9c5db199SXin Li    # enough to overwrite all the progress information we have sent, add some
484*9c5db199SXin Li    # more spaces.
485*9c5db199SXin Li    if final and self._update_len:
486*9c5db199SXin Li      print(' ' * self._update_len, '\r', end='')
487*9c5db199SXin Li
488*9c5db199SXin Li    self._pending_nl = -1
489*9c5db199SXin Li
490*9c5db199SXin Li  def _CheckStreamAndColor(self, stream, display):
491*9c5db199SXin Li    """Check that we're writing to the same stream as last call.  No?  New line.
492*9c5db199SXin Li
493*9c5db199SXin Li    If starting a new line, set the color correctly:
494*9c5db199SXin Li      stdout  Magenta
495*9c5db199SXin Li      stderr  Red
496*9c5db199SXin Li      other   White / no colors
497*9c5db199SXin Li
498*9c5db199SXin Li    Args:
499*9c5db199SXin Li      stream: The stream we're going to write to.
500*9c5db199SXin Li      display: True to display it on terms, False to suppress it.
501*9c5db199SXin Li    """
502*9c5db199SXin Li    if self._column > 0 and stream != self._cur_stream:
503*9c5db199SXin Li      self._FinishLine(display)
504*9c5db199SXin Li      if display:
505*9c5db199SXin Li        print(self._color.Stop())
506*9c5db199SXin Li
507*9c5db199SXin Li      self._column = 0
508*9c5db199SXin Li      self._line = ''
509*9c5db199SXin Li
510*9c5db199SXin Li    # Use colors for child output.
511*9c5db199SXin Li    if self._column == 0:
512*9c5db199SXin Li      self._FinishLine(display)
513*9c5db199SXin Li      if display:
514*9c5db199SXin Li        color = None
515*9c5db199SXin Li        if stream == sys.stdout:
516*9c5db199SXin Li          color = self._color.MAGENTA
517*9c5db199SXin Li        elif stream == sys.stderr:
518*9c5db199SXin Li          color = self._color.RED
519*9c5db199SXin Li        if color:
520*9c5db199SXin Li          sys.stdout.write(self._color.Start(color))
521*9c5db199SXin Li
522*9c5db199SXin Li      self._cur_stream = stream
523*9c5db199SXin Li
524*9c5db199SXin Li  def _Out(self, stream, text, display, newline=False, do_output_filter=True):
525*9c5db199SXin Li    """Output some text received from a child, or generated internally.
526*9c5db199SXin Li
527*9c5db199SXin Li    This method is the guts of the Operation class since it understands how to
528*9c5db199SXin Li    convert a series of output requests on different streams into something
529*9c5db199SXin Li    coherent for the user.
530*9c5db199SXin Li
531*9c5db199SXin Li    If the stream has changed, then a new line is started even if we were
532*9c5db199SXin Li    still halfway through the previous line. This prevents stdout and stderr
533*9c5db199SXin Li    becoming mixed up quite so badly.
534*9c5db199SXin Li
535*9c5db199SXin Li    We use color to indicate lines which are stdout and stderr. If the output
536*9c5db199SXin Li    received from the child has color codes in it already, we pass these
537*9c5db199SXin Li    through, so our colors can be overridden. If output is redirected then we
538*9c5db199SXin Li    do not add color by default. Note that nothing stops the child from adding
539*9c5db199SXin Li    it, but since we present ourselves as a terminal to the child, one might
540*9c5db199SXin Li    hope that the child will not generate color.
541*9c5db199SXin Li
542*9c5db199SXin Li    If display is False, then we will not actually send this text to the
543*9c5db199SXin Li    terminal. This is uses when verbose is required to be False.
544*9c5db199SXin Li
545*9c5db199SXin Li    Args:
546*9c5db199SXin Li      stream: stream on which the text was received:
547*9c5db199SXin Li        sys.stdout    - received on stdout
548*9c5db199SXin Li        sys.stderr    - received on stderr
549*9c5db199SXin Li        None          - generated by us / internally
550*9c5db199SXin Li      text: text to output
551*9c5db199SXin Li      display: True to display it on terms, False to suppress it
552*9c5db199SXin Li      newline: True to start a new line after this text, False to put the next
553*9c5db199SXin Li        lot of output immediately after this.
554*9c5db199SXin Li      do_output_filter: True to look through output for errors and progress.
555*9c5db199SXin Li    """
556*9c5db199SXin Li    self._CheckStreamAndColor(stream, display)
557*9c5db199SXin Li
558*9c5db199SXin Li    # Output what we have, and remember what column we are up to.
559*9c5db199SXin Li    if display:
560*9c5db199SXin Li      sys.stdout.write(text)
561*9c5db199SXin Li      self._column += len(text)
562*9c5db199SXin Li      # If a newline is required, remember to output it later.
563*9c5db199SXin Li      if newline:
564*9c5db199SXin Li        self._pending_nl = self._column
565*9c5db199SXin Li        self._column = 0
566*9c5db199SXin Li
567*9c5db199SXin Li    self._line += text
568*9c5db199SXin Li
569*9c5db199SXin Li    # If we now have a whole line, check it for errors and progress.
570*9c5db199SXin Li    if newline:
571*9c5db199SXin Li      if do_output_filter:
572*9c5db199SXin Li        self._FilterOutputForErrors(self._line, print_error=not display)
573*9c5db199SXin Li        self._FilterOutputForProgress(self._line)
574*9c5db199SXin Li      self._line = ''
575*9c5db199SXin Li
576*9c5db199SXin Li  def Output(self, stream, data):
577*9c5db199SXin Li    r"""Handle the output of a block of text from the subprocess.
578*9c5db199SXin Li
579*9c5db199SXin Li    All subprocess output should be sent through this method. It is split into
580*9c5db199SXin Li    lines which are processed separately using the _Out() method.
581*9c5db199SXin Li
582*9c5db199SXin Li    Args:
583*9c5db199SXin Li      stream: Which file the output come in on:
584*9c5db199SXin Li        sys.stdout: stdout
585*9c5db199SXin Li        sys.stderr: stderr
586*9c5db199SXin Li        None: Our own internal output
587*9c5db199SXin Li      data: Output data as a big string, potentially containing many lines of
588*9c5db199SXin Li        text. Each line should end with \r\n. There is no requirement to send
589*9c5db199SXin Li        whole lines - this method happily handles fragments and tries to
590*9c5db199SXin Li        present then to the user as early as possible
591*9c5db199SXin Li
592*9c5db199SXin Li    #TODO(sjg): Just use a list as the input parameter to avoid the split.
593*9c5db199SXin Li    """
594*9c5db199SXin Li    # We cannot use splitlines() here as we need this exact behavior
595*9c5db199SXin Li    lines = data.split('\r\n')
596*9c5db199SXin Li
597*9c5db199SXin Li    # Output each full line, with a \n after it.
598*9c5db199SXin Li    for line in lines[:-1]:
599*9c5db199SXin Li      self._Out(stream, line, display=self.verbose, newline=True)
600*9c5db199SXin Li
601*9c5db199SXin Li    # If we have a partial line at the end, output what we have.
602*9c5db199SXin Li    # We will continue it later.
603*9c5db199SXin Li    if lines[-1]:
604*9c5db199SXin Li      self._Out(stream, lines[-1], display=self.verbose)
605*9c5db199SXin Li
606*9c5db199SXin Li    # Flush so that the terminal will receive partial line output (now!)
607*9c5db199SXin Li    sys.stdout.flush()
608*9c5db199SXin Li
609*9c5db199SXin Li  def Outline(self, line):
610*9c5db199SXin Li    r"""Output a line of text to the display.
611*9c5db199SXin Li
612*9c5db199SXin Li    This outputs text generated internally, such as a warning message or error
613*9c5db199SXin Li    summary. It ensures that our message plays nicely with child output if
614*9c5db199SXin Li    any.
615*9c5db199SXin Li
616*9c5db199SXin Li    Args:
617*9c5db199SXin Li      line: text to output (without \n on the end)
618*9c5db199SXin Li    """
619*9c5db199SXin Li    self._Out(None, line, display=True, newline=True)
620*9c5db199SXin Li    self._FinishLine(display=True)
621*9c5db199SXin Li
622*9c5db199SXin Li  def Info(self, line):
623*9c5db199SXin Li    r"""Output a line of information text to the display in verbose mode.
624*9c5db199SXin Li
625*9c5db199SXin Li    Args:
626*9c5db199SXin Li      line: text to output (without \n on the end)
627*9c5db199SXin Li    """
628*9c5db199SXin Li    self._Out(None, self._color.Color(self._color.BLUE, line),
629*9c5db199SXin Li              display=self.verbose, newline=True, do_output_filter=False)
630*9c5db199SXin Li    self._FinishLine(display=True)
631*9c5db199SXin Li
632*9c5db199SXin Li  def Notice(self, line):
633*9c5db199SXin Li    r"""Output a line of notification text to the display.
634*9c5db199SXin Li
635*9c5db199SXin Li    Args:
636*9c5db199SXin Li      line: text to output (without \n on the end)
637*9c5db199SXin Li    """
638*9c5db199SXin Li    self._Out(None, self._color.Color(self._color.GREEN, line),
639*9c5db199SXin Li              display=True, newline=True, do_output_filter=False)
640*9c5db199SXin Li    self._FinishLine(display=True)
641*9c5db199SXin Li
642*9c5db199SXin Li  def Warning(self, line):
643*9c5db199SXin Li    r"""Output a line of warning text to the display.
644*9c5db199SXin Li
645*9c5db199SXin Li    Args:
646*9c5db199SXin Li      line: text to output (without \n on the end)
647*9c5db199SXin Li    """
648*9c5db199SXin Li    self._Out(None, self._color.Color(self._color.YELLOW, line),
649*9c5db199SXin Li              display=True, newline=True, do_output_filter=False)
650*9c5db199SXin Li    self._FinishLine(display=True)
651*9c5db199SXin Li
652*9c5db199SXin Li  def Error(self, line):
653*9c5db199SXin Li    r"""Output a line of error text to the display.
654*9c5db199SXin Li
655*9c5db199SXin Li    Args:
656*9c5db199SXin Li      line: text to output (without \n on the end)
657*9c5db199SXin Li    """
658*9c5db199SXin Li    self._Out(None, self._color.Color(self._color.RED, line),
659*9c5db199SXin Li              display=True, newline=True, do_output_filter=False)
660*9c5db199SXin Li    self._FinishLine(display=True)
661*9c5db199SXin Li
662*9c5db199SXin Li  def Die(self, line):
663*9c5db199SXin Li    r"""Output a line of error text to the display and die.
664*9c5db199SXin Li
665*9c5db199SXin Li    Args:
666*9c5db199SXin Li      line: text to output (without \n on the end)
667*9c5db199SXin Li    """
668*9c5db199SXin Li    self.Error(line)
669*9c5db199SXin Li    sys.exit(1)
670*9c5db199SXin Li
671*9c5db199SXin Li  @contextlib.contextmanager
672*9c5db199SXin Li  def RequestVerbose(self, request):
673*9c5db199SXin Li    """Perform something in verbose mode if the user hasn't disallowed it
674*9c5db199SXin Li
675*9c5db199SXin Li    This is intended to be used with something like:
676*9c5db199SXin Li
677*9c5db199SXin Li      with oper.RequestVerbose(True):
678*9c5db199SXin Li        ... do some things that generate output
679*9c5db199SXin Li
680*9c5db199SXin Li    Args:
681*9c5db199SXin Li      request: True to request verbose mode if available, False to do nothing.
682*9c5db199SXin Li    """
683*9c5db199SXin Li    old_verbose = self.verbose
684*9c5db199SXin Li    if request and not self.explicit_verbose:
685*9c5db199SXin Li      self.verbose = True
686*9c5db199SXin Li    try:
687*9c5db199SXin Li      yield
688*9c5db199SXin Li    finally:
689*9c5db199SXin Li      self.verbose = old_verbose
690