depot_tools: Simplify CheckCallAndFilter[AndHeader]

- Merge CheckCallAndFilter[AndHeader]
- print_stdout will cause command output to be redirected to sys.stdout.
- Made compatible with Python 3.

Bug: 984182
Change-Id: Ida30e295b872c8c1a1474a376a90aecea924f307
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/1727404
Commit-Queue: Edward Lesmes <ehmaldonado@chromium.org>
Reviewed-by: Robbie Iannucci <iannucci@chromium.org>
diff --git a/gclient_utils.py b/gclient_utils.py
index 3999fab..0957ec1 100644
--- a/gclient_utils.py
+++ b/gclient_utils.py
@@ -312,38 +312,6 @@
   return ' '.join(pipes.quote(arg) for arg in args)
 
 
-def CheckCallAndFilterAndHeader(args, always=False, header=None, **kwargs):
-  """Adds 'header' support to CheckCallAndFilter.
-
-  If |always| is True, a message indicating what is being done
-  is printed to stdout all the time even if not output is generated. Otherwise
-  the message header is printed only if the call generated any ouput.
-  """
-  stdout = kwargs.setdefault('stdout', sys.stdout)
-  if header is None:
-    # The automatically generated header only prepends newline if always is
-    # false: always is usually set to false if there's an external progress
-    # display, and it's better not to clobber it in that case.
-    header = "%s________ running '%s' in '%s'\n" % (
-                 '' if always else '\n',
-                 ' '.join(args), kwargs.get('cwd', '.'))
-
-  if always:
-    stdout.write(header)
-  else:
-    filter_fn = kwargs.get('filter_fn')
-    def filter_msg(line):
-      if line is None:
-        stdout.write(header)
-      elif filter_fn:
-        filter_fn(line)
-    kwargs['filter_fn'] = filter_msg
-    kwargs['call_filter_on_first_line'] = True
-  # Obviously.
-  kwargs.setdefault('print_stdout', True)
-  return CheckCallAndFilter(args, **kwargs)
-
-
 class Wrapper(object):
   """Wraps an object, acting as a transparent proxy for all properties by
   default.
@@ -355,22 +323,6 @@
     return getattr(self._wrapped, name)
 
 
-class WriteToStdout(Wrapper):
-  """Creates a file object clone to also print to sys.stdout."""
-  def __init__(self, wrapped):
-    super(WriteToStdout, self).__init__(wrapped)
-    if not hasattr(self, 'lock'):
-      self.lock = threading.Lock()
-
-  def write(self, out, *args, **kwargs):
-    self._wrapped.write(out, *args, **kwargs)
-    self.lock.acquire()
-    try:
-      sys.stdout.write(out, *args, **kwargs)
-    finally:
-      self.lock.release()
-
-
 class AutoFlush(Wrapper):
   """Creates a file object clone to automatically flush after N seconds."""
   def __init__(self, wrapped, delay):
@@ -536,9 +488,9 @@
           print('  ', zombie.pid, file=sys.stderr)
 
 
-def CheckCallAndFilter(args, stdout=None, filter_fn=None,
-                       print_stdout=None, call_filter_on_first_line=False,
-                       retry=False, **kwargs):
+def CheckCallAndFilter(args, print_stdout=False, filter_fn=None,
+                       show_header=False, always_show_header=False, retry=False,
+                       **kwargs):
   """Runs a command and calls back a filter function if needed.
 
   Accepts all subprocess2.Popen() parameters plus:
@@ -546,28 +498,68 @@
     filter_fn: A function taking a single string argument called with each line
                of the subprocess2's output. Each line has the trailing newline
                character trimmed.
-    stdout: Can be any bufferable output.
+    show_header: Whether to display a header before the command output.
+    always_show_header: Show header even when the command produced no output.
     retry: If the process exits non-zero, sleep for a brief interval and try
            again, up to RETRY_MAX times.
 
   stderr is always redirected to stdout.
+
+  Returns the output of the command as a binary string.
   """
-  assert print_stdout or filter_fn
-  stdout = stdout or sys.stdout
-  output = io.BytesIO()
-  filter_fn = filter_fn or (lambda x: None)
+  def show_header_if_necessary(needs_header, attempt):
+    """Show the header at most once."""
+    if not needs_header[0]:
+      return
+
+    needs_header[0] = False
+    # Automatically generated header. We only prepend a newline if
+    # always_show_header is false, since it usually indicates there's an
+    # external progress display, and it's better not to clobber it in that case.
+    header = '' if always_show_header else '\n'
+    header += '________ running \'%s\' in \'%s\'' % (
+                  ' '.join(args), kwargs.get('cwd', '.'))
+    if attempt:
+      header += ' attempt %s / %s' % (attempt + 1, RETRY_MAX + 1)
+    header += '\n'
+
+    if print_stdout:
+      sys.stdout.write(header)
+    if filter_fn:
+      filter_fn(header)
+
+  def filter_line(command_output, line_start):
+    """Extract the last line from command output and filter it."""
+    if not filter_fn or line_start is None:
+      return
+    command_output.seek(line_start)
+    filter_fn(command_output.read().decode('utf-8'))
+
+  # Initialize stdout writer if needed. On Python 3, sys.stdout does not accept
+  # byte inputs and sys.stdout.buffer must be used instead.
+  if print_stdout:
+    sys.stdout.flush()
+    stdout_write = getattr(sys.stdout, 'buffer', sys.stdout).write
+  else:
+    stdout_write = lambda _: None
 
   sleep_interval = RETRY_INITIAL_SLEEP
   run_cwd = kwargs.get('cwd', os.getcwd())
-  for _ in range(RETRY_MAX + 1):
+  for attempt in range(RETRY_MAX + 1):
     kid = subprocess2.Popen(
         args, bufsize=0, stdout=subprocess2.PIPE, stderr=subprocess2.STDOUT,
         **kwargs)
 
     GClientChildren.add(kid)
 
-    # Do a flush of stdout before we begin reading from the subprocess2's stdout
-    stdout.flush()
+    # Store the output of the command regardless of the value of print_stdout or
+    # filter_fn.
+    command_output = io.BytesIO()
+
+    # Passed as a list for "by ref" semantics.
+    needs_header = [show_header]
+    if always_show_header:
+      show_header_if_necessary(needs_header, attempt)
 
     # Also, we need to forward stdout to prevent weird re-ordering of output.
     # This has to be done on a per byte basis to make sure it is not buffered:
@@ -575,25 +567,29 @@
     # input, no end-of-line character is output after the prompt and it would
     # not show up.
     try:
-      in_byte = kid.stdout.read(1)
-      if in_byte:
-        if call_filter_on_first_line:
-          filter_fn(None)
-        in_line = b''
-        while in_byte:
-          output.write(in_byte)
-          if print_stdout:
-            stdout.write(in_byte)
-          if in_byte not in ['\r', '\n']:
-            in_line += in_byte
-          else:
-            filter_fn(in_line)
-            in_line = b''
-          in_byte = kid.stdout.read(1)
-        # Flush the rest of buffered output. This is only an issue with
-        # stdout/stderr not ending with a \n.
-        if len(in_line):
-          filter_fn(in_line)
+      line_start = None
+      while True:
+        in_byte = kid.stdout.read(1)
+        is_newline = in_byte in (b'\n', b'\r')
+        if not in_byte:
+          break
+
+        show_header_if_necessary(needs_header, attempt)
+
+        if is_newline:
+          filter_line(command_output, line_start)
+          line_start = None
+        elif line_start is None:
+          line_start = command_output.tell()
+
+        stdout_write(in_byte)
+        command_output.write(in_byte)
+
+      # Flush the rest of buffered output.
+      sys.stdout.flush()
+      if line_start is not None:
+        filter_line(command_output, line_start)
+
       rv = kid.wait()
       kid.stdout.close()
 
@@ -606,13 +602,16 @@
       raise
 
     if rv == 0:
-      return output.getvalue()
+      return command_output.getvalue()
+
     if not retry:
       break
+
     print("WARNING: subprocess '%s' in %s failed; will retry after a short "
           'nap...' % (' '.join('"%s"' % x for x in args), run_cwd))
     time.sleep(sleep_interval)
     sleep_interval *= 2
+
   raise subprocess2.CalledProcessError(
       rv, args, kwargs.get('cwd', None), None, None)
 
@@ -635,6 +634,7 @@
         The line will be skipped if predicate(line) returns False.
       out_fh: File handle to write output to.
     """
+    self.first_line = True
     self.last_time = 0
     self.time_throttle = time_throttle
     self.predicate = predicate
@@ -656,7 +656,9 @@
       elif now - self.last_time < self.time_throttle:
         return
     self.last_time = now
-    self.out_fh.write('[%s] ' % Elapsed())
+    if not self.first_line:
+      self.out_fh.write('[%s] ' % Elapsed())
+    self.first_line = False
     print(line, file=self.out_fh)