scripts/tools/annotee_indenter.py

#!/usr/bin/env python
# Copyright 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""annotee_indenter parses annotated input and indents each step level.

Used by ../slave/recipes/infra/try_recipe.py to execute sub-recipes
and indent output automatically.
"""

import argparse
import collections
import errno
import os
import select
import subprocess
import sys
import threading
import Queue


QUEUE_TERMINATE_ITEM = object()


def main(argv):
  parser = argparse.ArgumentParser(os.path.basename(__file__))
  parser.add_argument('--cwd', help='The cwd for executing child process.',
                      default=os.getcwd())
  parser.add_argument('--base-level', help='The base level for indents.',
                      type=int, required=True)
  parser.add_argument('--use-python-executable', action='store_true',
                      help='execute python command using same executbale, '
                           'as running this one.')
  parser.add_argument('sub', nargs='+')
  opts = parser.parse_args(argv)
  cmd = opts.sub
  if opts.use_python_executable:
    cmd = [sys.executable] + cmd
  return run(cmd, opts.cwd, opts.base_level, process_output)


def run(args, cwd, base_level, target):
  proc = subprocess.Popen(args, cwd=cwd,
                          stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  queue = Queue.Queue()
  t = threading.Thread(target=target, args=(queue, base_level))
  t.start()
  try:
    if sys.platform == "win32":
      communicate_win(proc, queue)
    else:
      communicate_posix(proc, queue)
  finally:
    # Ensure that on KeyboardInterrupt, the thread terminates gracefully.
    queue.put(QUEUE_TERMINATE_ITEM)
    t.join()
  return proc.wait()


def communicate_posix(proc, queue):
  """While process runs, puts tuples of (stream name, data) into queue."""

  def reader(rlist, stream_name):
    stream = getattr(proc, stream_name)
    if stream in rlist:
      data = os.read(stream.fileno(), 1024)
      if data == '':
        stream.close()
        read_set.remove(stream)
      queue.put((stream_name, data))

  read_set = [proc.stdout, proc.stderr]
  while read_set:
    try:
      rlist, _, _ = select.select(read_set, [], [])
    except select.error, e:
      if e.args[0] == errno.EINTR:
        continue
      raise
    reader(rlist, 'stdout')
    reader(rlist, 'stderr')


def communicate_win(proc, queue):
  def reader(stream_name):
    stream = getattr(proc, stream_name)
    while True:
      data = stream.read(1024)
      if data == '':
        break
      queue.put((stream_name, data))

  threads = []
  for stream_name in ['stdout', 'stderr']:
    t = threading.Thread(target=reader, args=(stream_name, ))
    t.daemon = True
    t.start()
    threads.append(t)

  for t in threads:
    t.join()


def line_generator(queue, stderr_stream=None):
  if stderr_stream is None:
    stderr_stream = sys.stderr
  # remainder stores the line which started, but not finished.
  remainder = {}
  while True:
    item = queue.get()
    if item == QUEUE_TERMINATE_ITEM:
      break
    stream, data = item
    data = remainder.pop(stream, '') + data
    if data == '':
      continue
    lines = data.splitlines(True)  # Keep endings.
    if not lines[-1].endswith('\n'):
      # Last line isn't finished yet, put into remainder.
      remainder[stream] = lines.pop()
    for l in lines:
      if stream == 'stderr':
        stderr_stream.write(l)
      else:
        yield l
  # Empty remainder, if any, at the end.
  if 'stderr' in remainder:
    stderr_stream.write(remainder.pop('stderr'))
  if 'stdout' in remainder:
    yield remainder.pop('stdout')
  assert not remainder


def process_output(queue, base_level):
  return indent(base_level, line_generator(queue), sys.stdout)


def indent(base_level, input_lines, output_stream):
  for l in input_lines:
    if l.startswith('@@@STEP_NEST_LEVEL@'):
      # Example: @@@STEP_NEST_LEVEL@1@@@
      level = int(l[len('@@@STEP_NEST_LEVEL@'):].strip().strip('@'))
      l = l.replace(str(level), str(level + base_level))
    output_stream.write(l)
    if l.startswith('@@@STEP_STARTED@@@'):
      output_stream.write('@@@STEP_NEST_LEVEL@%d@@@\n' % base_level)
  return 0


if __name__ == '__main__':
  sys.exit(main(sys.argv[1:]))