Support multiple try masters when sending tries to rietveld.

presubmit_support.py

 #!/usr/bin/env python
 # Copyright (c) 2012 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.
 
 """Enables directory-specific presubmit checks to run at upload and/or commit.
 """
 
 __version__ = '1.8.0'
 
 # TODO(joi) Add caching where appropriate/needed. The API is designed to allow
 # caching (between all different invocations of presubmit scripts for a given
 # change). We should add it as our presubmit scripts start feeling slow.
 
 import cpplint
 import cPickle  # Exposed through the API.
 import cStringIO  # Exposed through the API.
 import contextlib
 import fnmatch
 import glob
 import inspect
+import itertools
 import json  # Exposed through the API.
 import logging
 import marshal  # Exposed through the API.
 import multiprocessing
 import optparse
 import os  # Somewhat exposed through the API.
 import pickle  # Exposed through the API.
 import random
 import re  # Exposed through the API.
 import sys  # Parts exposed through API.
@@ skipping 968 matching lines @@
       results.append(p)
 
   logging.debug('Presubmit files: %s' % ','.join(results))
   return results
 
 
 class GetTrySlavesExecuter(object):
   @staticmethod
   def ExecPresubmitScript(script_text, presubmit_path, project, change):
     """Executes GetPreferredTrySlaves() from a single presubmit script.
+    
+    This will soon be deprecated and replaced by GetPreferredTryMasters().
 
     Args:
       script_text: The text of the presubmit script.
       presubmit_path: Project script to run.
       project: Project name to pass to presubmit script for bot selection.
 
     Return:
       A list of try slaves.
     """
     context = {}
@@ skipping 47 matching lines @@
              )
 
     # Ensure it's either all old-style or all new-style.
     if not valid_oldstyle(result) and not valid_newstyle(result):
       raise PresubmitFailure(
           'PRESUBMIT.py returned invalid trybot specification!')
 
     return result
 
 
+class GetTryMastersExecuter(object):
+  @staticmethod
+  def ExecPresubmitScript(script_text, presubmit_path, project, change):
+    """Executes GetPreferredTryMasters() from a single presubmit script.
+
+    Args:
+      script_text: The text of the presubmit script.
+      presubmit_path: Project script to run.
+      project: Project name to pass to presubmit script for bot selection.
+
+    Return:
+      A map of try masters to map of builders to set of tests.
+    """
+    context = {}
+    try:
+      exec script_text in context
+    except Exception, e:
+      raise PresubmitFailure('"%s" had an exception.\n%s'
+                             % (presubmit_path, e))
+
+    function_name = 'GetPreferredTryMasters'
+    if function_name not in context:
+      return {}
+    get_preferred_try_masters = context[function_name]
+    if not len(inspect.getargspec(get_preferred_try_masters)[0]) == 2:
+      raise PresubmitFailure(
+          'Expected function "GetPreferredTryMasters" to take two arguments.')
+    return get_preferred_try_masters(project, change)
+
+
 def DoGetTrySlaves(change,
                    changed_files,
                    repository_root,
                    default_presubmit,
                    project,
                    verbose,
                    output_stream):
-  """Get the list of try servers from the presubmit scripts.
+  """Get the list of try servers from the presubmit scripts (deprecated).
 
   Args:
     changed_files: List of modified files.
     repository_root: The repository root.
     default_presubmit: A default presubmit script to execute in any case.
     project: Optional name of a project used in selecting trybots.
     verbose: Prints debug info.
     output_stream: A stream to write debug output to.
 
   Return:
@@ skipping 30 matching lines @@
   slaves = list(slave_dict.items())
 
   slaves.extend(set(old_style))
 
   if slaves and verbose:
     output_stream.write(', '.join((str(x) for x in slaves)))
     output_stream.write('\n')
   return slaves
 
 
+def _MergeMasters(masters1, masters2):
+  """Merges two master maps. Merges also the tests of each builder."""
+  result = {}
+  for (master, builders) in itertools.chain(masters1.iteritems(),
+                                            masters2.iteritems()):
+    new_builders = result.setdefault(master, {})
+    for (builder, tests) in builders.iteritems():
+      new_builders.setdefault(builder, set([])).update(tests)
+  return result
+
+
+def DoGetTryMasters(change,
+                    changed_files,
+                    repository_root,
+                    default_presubmit,
+                    project,
+                    verbose,
+                    output_stream):
+  """Get the list of try masters from the presubmit scripts.
+
+  Args:
+    changed_files: List of modified files.
+    repository_root: The repository root.
+    default_presubmit: A default presubmit script to execute in any case.
+    project: Optional name of a project used in selecting trybots.
+    verbose: Prints debug info.
+    output_stream: A stream to write debug output to.
+
+  Return:
+    Map of try masters to map of builders to set of tests.
+  """
+  presubmit_files = ListRelevantPresubmitFiles(changed_files, repository_root)
+  if not presubmit_files and verbose:
+    output_stream.write("Warning, no PRESUBMIT.py found.\n")
+  results = {}
+  executer = GetTryMastersExecuter()
+
+  if default_presubmit:
+    if verbose:
+      output_stream.write("Running default presubmit script.\n")
+    fake_path = os.path.join(repository_root, 'PRESUBMIT.py')
+    results = _MergeMasters(results, executer.ExecPresubmitScript(
+        default_presubmit, fake_path, project, change))
+  for filename in presubmit_files:
+    filename = os.path.abspath(filename)
+    if verbose:
+      output_stream.write("Running %s\n" % filename)
+    # Accept CRLF presubmit script.
+    presubmit_script = gclient_utils.FileRead(filename, 'rU')
+    results = _MergeMasters(results, executer.ExecPresubmitScript(
+        presubmit_script, filename, project, change))
+
+  # Make sets to lists again for later JSON serialization.
+  for builders in results.itervalues():
+    for builder in builders:
+      builders[builder] = list(builders[builder])
+
+  if results and verbose:
+    output_stream.write('%s\n' % str(results))
+  return results
+
+
 class PresubmitExecuter(object):
   def __init__(self, change, committing, rietveld_obj, verbose):
     """
     Args:
       change: The Change object.
       committing: True if 'gcl commit' is running, False if 'gcl upload' is.
       rietveld_obj: rietveld.Rietveld client object.
     """
     self.change = change
     self.committing = committing
@@ skipping 363 matching lines @@
   except PresubmitFailure, e:
     print >> sys.stderr, e
     print >> sys.stderr, 'Maybe your depot_tools is out of date?'
     print >> sys.stderr, 'If all fails, contact maruel@'
     return 2
 
 
 if __name__ == '__main__':
   fix_encoding.fix_encoding()
   sys.exit(Main(None))