git_map.py

#!/usr/bin/env python3
# Copyright 2014 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.

"""
usage: git map [-h] [--help] [<args>]

Enhances `git log --graph` view with information on commit branches + tags that
point to them. Items are colorized as follows:

  * Cyan    - Currently checked out branch
  * Green   - Local branch
  * Red     - Remote branches
  * Magenta - Tags
  * White   - Merge Base Markers
  * Blue background - The currently checked out commit
"""

from __future__ import unicode_literals

import os
import sys

import git_common
import setup_color
import subprocess2

from third_party import colorama


if sys.version_info.major == 2:
  # On Python 3, BrokenPipeError is raised instead.
  # pylint:disable=redefined-builtin
  BrokenPipeError = IOError


RESET = colorama.Fore.RESET + colorama.Back.RESET + colorama.Style.RESET_ALL
BRIGHT = colorama.Style.BRIGHT

BLUE_BACK = colorama.Back.BLUE + BRIGHT
BRIGHT_RED = colorama.Fore.RED + BRIGHT
CYAN = colorama.Fore.CYAN + BRIGHT
GREEN = colorama.Fore.GREEN + BRIGHT
MAGENTA = colorama.Fore.MAGENTA + BRIGHT
RED = colorama.Fore.RED
WHITE = colorama.Fore.WHITE + BRIGHT
YELLOW = colorama.Fore.YELLOW


def _print_help(outbuf):
  names = {
    'Cyan': CYAN,
    'Green': GREEN,
    'Magenta': MAGENTA,
    'Red': RED,
    'White': WHITE,
    'Blue background': BLUE_BACK,
  }
  msg = ''
  for line in __doc__.splitlines():
    for name, color in names.items():
      if name in line:
        msg += line.replace('* ' + name, color + '* ' + name + RESET) + '\n'
        break
    else:
      msg += line + '\n'
  outbuf.write(msg.encode('utf-8', 'replace'))


def _color_branch(branch, all_branches, all_tags, current):
  if branch in (current, 'HEAD -> ' + current):
    color = CYAN
    current = None
  elif branch in all_branches:
    color = GREEN
    all_branches.remove(branch)
  elif branch in all_tags:
    color = MAGENTA
  elif branch.startswith('tag: '):
    color = MAGENTA
    branch = branch[len('tag: '):]
  else:
    color = RED
  return color + branch + RESET


def _color_branch_list(branch_list, all_branches, all_tags, current):
  if not branch_list:
    return ''
  colored_branches = (GREEN + ', ').join(
      _color_branch(branch, all_branches, all_tags, current)
      for branch in branch_list if branch != 'HEAD')
  return (GREEN + '(' + colored_branches + GREEN + ') ' + RESET)


def _parse_log_line(line):
  graph, branch_list, commit_date, subject = (
      line.decode('utf-8', 'replace').strip().split('\x00'))
  branch_list = [] if not branch_list else branch_list.split(', ')
  commit = graph.split()[-1]
  graph = graph[:-len(commit)]
  return graph, commit, branch_list, commit_date, subject


def main(argv, outbuf):
  if '-h' in argv or '--help' in argv:
    _print_help(outbuf)
    return 0

  map_extra = git_common.get_config_list('depot_tools.map_extra')
  cmd = [
      git_common.GIT_EXE, 'log', git_common.root(),
      '--graph', '--branches', '--tags', '--color=always', '--date=short',
      '--pretty=format:%H%x00%D%x00%cd%x00%s'
  ] + map_extra + argv

  log_proc = subprocess2.Popen(cmd, stdout=subprocess2.PIPE, shell=False)

  current = git_common.current_branch()
  all_tags = set(git_common.tags())
  all_branches = set(git_common.branches())
  if current in all_branches:
    all_branches.remove(current)

  merge_base_map = {}
  for branch in all_branches:
    merge_base = git_common.get_or_create_merge_base(branch)
    if merge_base:
      merge_base_map.setdefault(merge_base, set()).add(branch)

  for merge_base, branches in merge_base_map.items():
    merge_base_map[merge_base] = ', '.join(branches)

  try:
    for line in log_proc.stdout:
      if b'\x00' not in line:
        outbuf.write(line)
        continue

      graph, commit, branch_list, commit_date, subject = _parse_log_line(line)

      if 'HEAD' in branch_list:
        graph = graph.replace('*', BLUE_BACK + '*')

      line = '{graph}{commit}\t{branches}{date} ~ {subject}'.format(
          graph=graph,
          commit=BRIGHT_RED + commit[:10] + RESET,
          branches=_color_branch_list(
              branch_list, all_branches, all_tags, current),
          date=YELLOW + commit_date + RESET,
          subject=subject)

      if commit in merge_base_map:
        line += '    <({})'.format(WHITE + merge_base_map[commit] + RESET)

      line += os.linesep
      outbuf.write(line.encode('utf-8', 'replace'))
  except (BrokenPipeError, KeyboardInterrupt):
    pass
  return 0


if __name__ == '__main__':
  setup_color.init()
  with git_common.less() as less_input:
    sys.exit(main(sys.argv[1:], less_input))