Give git map-branches extra information behind -v and -vv flags.

git_map_branches.py

 #!/usr/bin/env python
 # 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.
 
-"""
-Provides a short mapping of all the branches in your local repo, organized by
-their upstream ('tracking branch') layout. Example:
+"""Provides a short mapping of all the branches in your local repo, organized
+by their upstream ('tracking branch') layout.
 
+Example:
 origin/master
   cool_feature
     dependent_feature
     other_dependent_feature
   other_feature
 
 Branches are colorized as follows:
   * Red - a remote branch (usually the root of all local branches)
   * Cyan - a local branch which is the same as HEAD
     * Note that multiple branches may be Cyan, if they are all on the same
       commit, and you have that commit checked out.
   * Green - a local branch
   * Magenta - a tag
   * Magenta '{NO UPSTREAM}' - If you have local branches which do not track any
     upstream, then you will see this.
 """
 
+import argparse
 import collections
 import sys
 
 from third_party import colorama
 from third_party.colorama import Fore, Style
 
-from git_common import current_branch, branches, upstream, hash_one, hash_multi
-from git_common import tags
+from git_common import current_branch, upstream, hash_one
+from git_common import tags, get_all_tracking_info, normalized_version
+from git_common import MIN_UPSTREAM_TRACK_GIT_VERSION
+
+import git_cl
 
 NO_UPSTREAM = '{NO UPSTREAM}'
-
-def color_for_branch(branch, branch_hash, cur_hash, tag_set):
-  if branch.startswith('origin'):
-    color = Fore.RED
-  elif branch == NO_UPSTREAM or branch in tag_set:
-    color = Fore.MAGENTA
-  elif branch_hash == cur_hash:
-    color = Fore.CYAN
-  else:
-    color = Fore.GREEN
-
-  if branch_hash == cur_hash:
-    color += Style.BRIGHT
-  else:
-    color += Style.NORMAL
-
-  return color
-
-
-def print_branch(cur, cur_hash, branch, branch_hashes, par_map, branch_map,
-                 tag_set, depth=0):
-  branch_hash = branch_hashes[branch]
-
-  color = color_for_branch(branch, branch_hash, cur_hash, tag_set)
-
-  suffix = ''
-  if cur == 'HEAD':
-    if branch_hash == cur_hash:
+DEFAULT_SEPARATOR = ' ' * 4
+
+
+class OutputManager(object):
+  """A class that manages a number of OutputLines and formats them into
+  aligned columns."""

[Matt Giuca, 2014/09/01 05:12:50]

The first sentence is supposed to fit on one line. If you need more detail, have
a blank line and then another paragraph (like a git commit).

[calamity, 2014/09/01 06:53:42]

Done.

+
+  def __init__(self):
+    self.lines = []
+    self.nocolor = False
+    self.max_column_lengths = []
+    self.num_columns = None
+
+  def append(self, line):
+    # All lines must have the same number of columns.
+    if not self.num_columns:
+      self.num_columns = len(line.columns)
+      self.max_column_lengths = [0] * self.num_columns
+    assert self.num_columns == len(line.columns)
+
+    if self.nocolor:
+      line.colors = [''] * self.num_columns
+
+    self.lines.append(line)
+
+    # Update maximum column lengths

[Matt Giuca, 2014/09/01 05:12:50]

Full stop.

[calamity, 2014/09/01 06:53:42]

Done.

+    for i, col in enumerate(line.columns):
+      self.max_column_lengths[i] = max(self.max_column_lengths[i], len(col))
+
+  def __str__(self):

[Matt Giuca, 2014/09/01 05:12:51]

I wouldn't use __str__ for this. __str__ should be short and not have newlines
(so that it is appropriate for putting in a print with "%s"). Call this method
as_string or something.

[calamity, 2014/09/01 06:53:42]

Done.

+    return '\n'.join(
+        l.as_padded_string(self.max_column_lengths) for l in self.lines)
+
+
+class OutputLine(object):
+  """A single line of data, consisting of an equal number of columns, colors and
+  separators."""

[Matt Giuca, 2014/09/01 05:12:50]

Same.

[calamity, 2014/09/01 06:53:42]

Done.

+
+  def __init__(self):
+    self.columns = []
+    self.separators = []
+    self.colors = []
+
+  def append(self, data, separator=DEFAULT_SEPARATOR, color=Fore.WHITE):
+    self.columns.append(data)
+    self.separators.append(separator)
+    self.colors.append(color)
+
+  def as_padded_string(self, max_column_lengths):
+    """"Returns the data as a string with each column padded to
+    |max_column_lengths|."""
+    output_string = ''
+    for i, (color, data, separator) in enumerate(
+        zip(self.colors, self.columns, self.separators)):
+      if max_column_lengths[i] == 0:
+        continue
+
+      padding = (max_column_lengths[i] - len(data)) * ' '
+      output_string += color + data + padding + separator
+
+    return output_string.rstrip()
+
+
+class BranchMapper(object):
+  """A class which constructs output representing the tree's branch structure"""

[Matt Giuca, 2014/09/01 05:12:51]

nit: Need a full stop at the end.

Also it needs to fit on one line: have fun!

I think you need documentation about what the major attributes of this class
represent. In particular, give some details about |tracking_info| -- what its
keys and values are.

[calamity, 2014/09/01 06:53:42]

Done.

+
+  def __init__(self):
+    self.verbosity = 0
+    self.output = OutputManager()
+    self.tracking_info = get_all_tracking_info()
+    self.__gone_branches = set()
+
+    # A map of parents to a list of their children.
+    self.parent_map = collections.defaultdict(list)
+    for branch in self.tracking_info.keys():

[Matt Giuca, 2014/09/01 05:12:51]

nit: Don't need .keys() (just "for branch in self.tracking_info").

Alternatively, do:
for branch, branch_info in self.tracking_info.iteritems()

and then |branch_info| is equal to self.tracking_info[branch].

[calamity, 2014/09/01 06:53:42]

Done.

+      parent = self.tracking_info[branch].get('upstream', None)
+      if not parent:
+        parent = NO_UPSTREAM
+      elif parent not in self.tracking_info:
+        branch_upstream = upstream(branch)
+        # If git can't find the upstream, mark the upstream as gone.
+        if branch_upstream:
+          parent = branch_upstream
+        else:
+          self.__gone_branches.add(parent)
+      self.parent_map[parent].append(branch)
+
+    self.current_branch = current_branch()
+    self.current_hash = self.tracking_info[self.current_branch]['hash']
+
+  def start(self):
+    tag_set = tags()
+
+    while self.parent_map:
+      for parent in sorted(self.parent_map.keys()):

[Matt Giuca, 2014/09/01 05:12:50]

Don't need .keys().

(Note that iterating over a dict produces its keys, so .keys() is never
necessary unless you actually directly want a list of keys.)

[calamity, 2014/09/01 06:53:42]

Done.

+        if parent in self.tracking_info:
+          continue
+
+        self.tracking_info[parent] = {
+            'hash': '' if self.__is_invalid_parent(parent)
+                       else hash_one(parent, short=True)}

[Matt Giuca, 2014/09/01 05:12:50]

Super hard to read. Break this expression out and indent it properly. Also I
can't find a citation for this, but I believe the correct Python style is to
have the "normal" case first and the "error" case second. So invert the logic
here.

[calamity, 2014/09/01 06:53:42]

N/A. But I fixed this to be a lot nicer =D

+
+        self.__append_branch(parent, tag_set)
+        break
+
+  def __is_invalid_parent(self, parent):
+    return parent == NO_UPSTREAM or parent in self.__gone_branches
+
+  def __color_for_branch(self, branch, branch_hash, tag_set):
+    if branch.startswith('origin'):
+      color = Fore.RED
+    elif self.__is_invalid_parent(branch) or branch in tag_set:
+      color = Fore.MAGENTA
+    elif branch_hash == self.current_hash:
+      color = Fore.CYAN
+    else:
+      color = Fore.GREEN
+
+    if branch_hash == self.current_hash:
+      color += Style.BRIGHT
+    else:
+      color += Style.NORMAL
+
+    return color
+
+  def __append_branch(self, branch, tag_set, depth=0):
+    """Recurses through the tree structure and appends an OutputLine to the
+    OutputManager for each branch."""
+    branch_hash = self.tracking_info[branch]['hash']
+
+    line = OutputLine()
+
+    # The branch name with appropriate indentation.
+    suffix = ''
+    if self.current_branch == 'HEAD':
+      if branch_hash == self.current_hash:
+        suffix = ' *'
+    elif branch == self.current_branch:
       suffix = ' *'

[Matt Giuca, 2014/09/01 05:12:50]

optional (might look weird): Combine this into one expression:
if branch == self.current_branch or (self.current_branch == 'HEAD' and
branch_hash == self.current_hash):

[calamity, 2014/09/01 06:53:42]

Done.

-  elif branch == cur:
-    suffix = ' *'
-
-  print color + "  "*depth + branch + suffix
-  for child in par_map.pop(branch, ()):
-    print_branch(cur, cur_hash, child, branch_hashes, par_map, branch_map,
-                 tag_set, depth=depth+1)
+    branch_string = (
+        '{%s:GONE}' % branch if branch in self.__gone_branches else branch)
+    main_string = '  ' * depth + branch_string + suffix
+    line.append(
+        main_string,
+        color=self.__color_for_branch(branch, branch_hash, tag_set))
+
+    # The branch hash.
+    if self.verbosity >= 2:
+      line.append(branch_hash, separator=' ', color=Fore.RED)
+
+    # The branch tracking status.
+    if self.verbosity >= 1:
+      ahead_string = ''
+      behind_string = ''
+      front_separator = ''
+      center_separator = ''
+      back_separator = ''
+      if branch in self.tracking_info and not self.__is_invalid_parent(
+          self.tracking_info[branch].get('upstream', NO_UPSTREAM)):
+        ahead_string = self.tracking_info[branch].get('ahead', '')
+        behind_string = self.tracking_info[branch].get('behind', '')
+
+        if ahead_string or behind_string:
+          front_separator = '['
+          back_separator = ']'
+
+        if ahead_string and behind_string:
+          center_separator = '|'
+
+      line.append(front_separator, separator=' ')
+      line.append(ahead_string, separator=' ', color=Fore.MAGENTA)
+      line.append(center_separator, separator=' ')
+      line.append(behind_string, separator=' ', color=Fore.MAGENTA)
+      line.append(back_separator)
+
+    # The rietveld issue associated with the branch.

[Matt Giuca, 2014/09/01 05:12:51]

Capital R.

[calamity, 2014/09/01 06:53:42]

Done.

+    if self.verbosity >= 2:
+      none_text = '' if self.__is_invalid_parent(branch) else 'None'
+      url = git_cl.Changelist(branchref=branch).GetIssueURL()
+      line.append(url or none_text, color=Fore.BLUE if url else Fore.WHITE)
+
+    self.output.append(line)
+
+    for child in sorted(self.parent_map.pop(branch, ())):
+      self.__append_branch(child, tag_set, depth=depth + 1)
 
 
 def main(argv):
   colorama.init()
-  assert len(argv) == 1, "No arguments expected"
-  branch_map = {}
-  par_map = collections.defaultdict(list)
-  for branch in branches():
-    par = upstream(branch) or NO_UPSTREAM
-    branch_map[branch] = par
-    par_map[par].append(branch)
-
-  current = current_branch()
-  hashes = hash_multi(current, *branch_map.keys())
-  current_hash = hashes[0]
-  par_hashes = {k: hashes[i+1] for i, k in enumerate(branch_map.iterkeys())}
-  par_hashes[NO_UPSTREAM] = 0
-  tag_set = tags()
-  while par_map:
-    for parent in par_map:
-      if parent not in branch_map:
-        if parent not in par_hashes:
-          par_hashes[parent] = hash_one(parent)
-        print_branch(current, current_hash, parent, par_hashes, par_map,
-                     branch_map, tag_set)
-        break
-
+  if normalized_version() < MIN_UPSTREAM_TRACK_GIT_VERSION:
+    print >> sys.stderr, (
+        'This tool will not show all tracking information for git version '
+        'earlier than ' +
+        '.'.join(str(x) for x in MIN_UPSTREAM_TRACK_GIT_VERSION) +
+        '. Please consider upgrading.')
+
+  parser = argparse.ArgumentParser(
+      description='Print a a tree of all branches parented by their upstreams')
+  parser.add_argument('-v', action='store_true',

[iannucci, 2014/08/29 18:41:47]

let's use the action 'count' for this, instead of having -v and -vv

[calamity, 2014/09/01 06:53:42]

Done.

+                      help='Display upstream tracking status')

[iannucci, 2014/08/29 18:41:47]

should we just do this by default? is there any reason not to?

[calamity, 2014/09/01 06:53:42]

Eh. Not really. I just didn't want to modify the original default functionality.

[Matt Giuca, 2014/09/02 23:43:33]

Seems consistent with git branch's behaviour of not being very helpful by
default :)

Seriously, I'd say leave it off by default for now, send out a PSA on
chromium-dev after it lands and see what people think, if they find issues with
it, etc. Then if enough people like it and after any bugs are fixed, turn it on
by default.

[calamity, 2014/09/03 01:43:21]

Done.

+  parser.add_argument(
+      '-vv',
+      action='store_true',
+      help='Like -v but also show branch hash and Rietveld URL.')
+  parser.add_argument('--no-color', action='store_true', dest='nocolor',
+                      help='Turn off colors.')
+
+  opts = parser.parse_args(argv[1:])
+
+  verbosity = 0
+  if opts.v:
+    verbosity = 1
+  if opts.vv:
+    verbosity = 2

[iannucci, 2014/08/29 18:41:47]

then you don't need this

[calamity, 2014/09/01 06:53:42]

Done.

+
+  mapper = BranchMapper()
+  mapper.verbosity = verbosity
+  mapper.output.nocolor = opts.nocolor
+  mapper.start()
+  print mapper.output

[iannucci, 2014/08/29 18:41:47]

not sure if this is clearer than just returning the output from start()

[calamity, 2014/09/01 06:53:42]

I think this is useful in case BranchMapper gets used for something else.
start() would be the data processing call after which clients could ask for
information.

 
 if __name__ == '__main__':
   sys.exit(main(sys.argv))
-