Make git-map-branches -vvv show CL status colors.

This CL makes git-map-branches show CL status colors like git cl status
when -vvv is used. Statuses are fetched in parallel for speed.

BUG=379849

Review URL: https://codereview.chromium.org/938583002

git-svn-id: svn://svn.chromium.org/chrome/trunk/tools/depot_tools@294414 0039d316-1c4b-4281-b951-d872f2087c98

git_cl.py

@@ -8,6 +8,7 @@
 """A git-command for integrating reviews on Rietveld."""
 
 from distutils.version import LooseVersion
+from multiprocessing.pool import ThreadPool
 import base64
 import glob
 import json
@@ -20,7 +21,6 @@ import stat
 import sys
 import tempfile
 import textwrap
-import threading
 import urllib2
 import urlparse
 import webbrowser
@@ -1342,6 +1342,51 @@ def color_for_status(status):
     'error': Fore.WHITE,
   }.get(status, Fore.WHITE)
 
+def fetch_cl_status(b):
+  """Fetches information for an issue and returns (branch, issue, color)."""
+  c = Changelist(branchref=b)
+  i = c.GetIssueURL()
+  status = c.GetStatus()
+  color = color_for_status(status)
+
+  if i and (not status or status == 'error'):
+    # The issue probably doesn't exist anymore.
+    i += ' (broken)'
+
+  return (b, i, color)
+
+def get_cl_statuses(branches, fine_grained, max_processes=None):
+  """Returns a blocking iterable of (branch, issue, color) for given branches.
+
+  If fine_grained is true, this will fetch CL statuses from the server.
+  Otherwise, simply indicate if there's a matching url for the given branches.
+
+  If max_processes is specified, it is used as the maximum number of processes
+  to spawn to fetch CL status from the server. Otherwise 1 process per branch is
+  spawned.
+  """
+  # Silence upload.py otherwise it becomes unwieldly.
+  upload.verbosity = 0
+
+  if fine_grained:
+    # Process one branch synchronously to work through authentication, then
+    # spawn processes to process all the other branches in parallel.
+    if branches:
+      yield fetch_cl_status(branches[0])
+
+      branches_to_fetch = branches[1:]
+      pool = ThreadPool(
+          min(max_processes, len(branches_to_fetch))
+              if max_processes is not None
+              else len(branches_to_fetch))
+      for x in pool.imap_unordered(fetch_cl_status, branches_to_fetch):
+        yield x
+  else:
+    # Do not use GetApprovingReviewers(), since it requires an HTTP request.
+    for b in branches:
+      c = Changelist(branchref=b)
+      url = c.GetIssueURL()
+      yield (b, url, Fore.BLUE if url else Fore.WHITE)
 
 def CMDstatus(parser, args):
   """Show status of changelists.
@@ -1360,6 +1405,9 @@ def CMDstatus(parser, args):
                     help='print only specific field (desc|id|patch|url)')
   parser.add_option('-f', '--fast', action='store_true',
                     help='Do not retrieve review status')
+  parser.add_option(
+      '-j', '--maxjobs', action='store', type=int,
+      help='The maximum number of jobs to use when retrieving review status')
   (options, args) = parser.parse_args(args)
   if args:
     parser.error('Unsupported args: %s' % args)
@@ -1391,49 +1439,17 @@ def CMDstatus(parser, args):
   branches = [c.GetBranch() for c in changes]
   alignment = max(5, max(len(b) for b in branches))
   print 'Branches associated with reviews:'
-  # Adhoc thread pool to request data concurrently.
-  output = Queue.Queue()
-
-  # Silence upload.py otherwise it becomes unweldly.
-  upload.verbosity = 0
-
-  if not options.fast:
-    def fetch(b):
-      """Fetches information for an issue and returns (branch, issue, color)."""
-      c = Changelist(branchref=b)
-      i = c.GetIssueURL()
-      status = c.GetStatus()
-      color = color_for_status(status)
-
-      if i and (not status or status == 'error'):
-        # The issue probably doesn't exist anymore.
-        i += ' (broken)'
-
-      output.put((b, i, color))
-
-    # Process one branch synchronously to work through authentication, then
-    # spawn threads to process all the other branches in parallel.
-    if branches:
-      fetch(branches[0])
-    threads = [
-      threading.Thread(target=fetch, args=(b,)) for b in branches[1:]]
-    for t in threads:
-      t.daemon = True
-      t.start()
-  else:
-    # Do not use GetApprovingReviewers(), since it requires an HTTP request.
-    for b in branches:
-      c = Changelist(branchref=b)
-      url = c.GetIssueURL()
-      output.put((b, url, Fore.BLUE if url else Fore.WHITE))
+  output = get_cl_statuses(branches,
+                           fine_grained=not options.fast,
+                           max_processes=options.maxjobs)
 
-  tmp = {}
+  branch_statuses = {}
   alignment = max(5, max(len(ShortBranchName(b)) for b in branches))
   for branch in sorted(branches):
-    while branch not in tmp:
-      b, i, color = output.get()
-      tmp[b] = (i, color)
-    issue, color = tmp.pop(branch)
+    while branch not in branch_statuses:
+      b, i, color = output.next()
+      branch_statuses[b] = (i, color)
+    issue, color = branch_statuses.pop(branch)
     reset = Fore.RESET
     if not sys.stdout.isatty():
       color = ''

git_map_branches.py

@@ -109,6 +109,7 @@ class BranchMapper(object):
 
   def __init__(self):
     self.verbosity = 0
+    self.maxjobs = 0
     self.output = OutputManager()
     self.__gone_branches = set()
     self.__branches_info = None
@@ -116,10 +117,25 @@ class BranchMapper(object):
     self.__current_branch = None
     self.__current_hash = None
     self.__tag_set = None
+    self.__status_info = {}
 
   def start(self):
     self.__branches_info = get_branches_info(
         include_tracking_status=self.verbosity >= 1)
+    if (self.verbosity >= 2):
+      # Avoid heavy import unless necessary.
+      from git_cl import get_cl_statuses
+
+      status_info = get_cl_statuses(self.__branches_info.keys(),
+                                    fine_grained=self.verbosity > 2,
+                                    max_processes=self.maxjobs)
+
+      for _ in xrange(len(self.__branches_info)):
+        # This is a blocking get which waits for the remote CL status to be
+        # retrieved.
+        (branch, url, color) = status_info.next()
+        self.__status_info[branch] = (url, color);
+
     roots = set()
 
     # A map of parents to a list of their children.
@@ -238,11 +254,9 @@ class BranchMapper(object):
 
     # The Rietveld issue associated with the branch.
     if self.verbosity >= 2:
-      import git_cl  # avoid heavy import cost unless we need it
       none_text = '' if self.__is_invalid_parent(branch) else 'None'
-      url = git_cl.Changelist(
-          branchref=branch).GetIssueURL() if branch_hash else None
-      line.append(url or none_text, color=Fore.BLUE if url else Fore.WHITE)
+      (url, color) = self.__status_info[branch]
+      line.append(url or none_text, color=color)
 
     self.output.append(line)
 
@@ -265,12 +279,16 @@ def main(argv):
                       help='Display branch hash and Rietveld URL')
   parser.add_argument('--no-color', action='store_true', dest='nocolor',
                       help='Turn off colors.')
+  parser.add_argument(
+      '-j', '--maxjobs', action='store', type=int,
+      help='The number of jobs to use when retrieving review status')
 
   opts = parser.parse_args(argv)
 
   mapper = BranchMapper()
   mapper.verbosity = opts.v
   mapper.output.nocolor = opts.nocolor
+  mapper.maxjobs = opts.maxjobs
   mapper.start()
   print mapper.output.as_formatted_string()
   return 0