Add support for parallel presubmit unit testing.

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.6.2'
 
 # 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 collections
 import contextlib
 import fnmatch
 import glob
 import inspect
 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.
 import tempfile  # Exposed through the API.
 import time
 import traceback  # Exposed through the API.
 import types
@@ skipping 12 matching lines @@
 
 
 # Ask for feedback only once in program lifetime.
 _ASKED_FOR_FEEDBACK = False
 
 
 class PresubmitFailure(Exception):
   pass
 
 
+CommandData = collections.namedtuple('CommandData',
+                                     ['name', 'cmd', 'kwargs', 'message'])
+
 def normpath(path):
   '''Version of os.path.normpath that also changes backward slashes to
   forward slashes when not running on Windows.
   '''
   # This is safe to always do because the Windows version of os.path.normpath
   # will replace forward slashes with backward slashes.
   path = path.replace(os.sep, '/')
   return os.path.normpath(path)
 
 
@@ skipping 30 matching lines @@
 
   def write(self, s):
     self.written_output.append(s)
     if self.output_stream:
       self.output_stream.write(s)
 
   def getvalue(self):
     return ''.join(self.written_output)
 
 
+# Top level object so multiprocessing can pickle
+# Public access through OutputApi object.
+class _PresubmitResult(object):
+  """Base class for result objects."""
+  fatal = False
+  should_prompt = False
+
+  def __init__(self, message, items=None, long_text=''):
+    """
+    message: A short one-line message to indicate errors.
+    items: A list of short strings to indicate where errors occurred.
+    long_text: multi-line text output, e.g. from another tool
+    """
+    self._message = message
+    self._items = items or []
+    if items:
+      self._items = items
+    self._long_text = long_text.rstrip()
+
+  def handle(self, output):
+    output.write(self._message)
+    output.write('\n')
+    for index, item in enumerate(self._items):
+      output.write('  ')
+      # Write separately in case it's unicode.
+      output.write(str(item))
+      if index < len(self._items) - 1:
+        output.write(' \\')
+      output.write('\n')
+    if self._long_text:
+      output.write('\n***************\n')
+      # Write separately in case it's unicode.
+      output.write(self._long_text)
+      output.write('\n***************\n')
+    if self.fatal:
+      output.fail()
+
+
+# Top level object so multiprocessing can pickle
+# Public access through OutputApi object.
+class _PresubmitAddReviewers(_PresubmitResult):
+  """Add some suggested reviewers to the change."""
+  def __init__(self, reviewers):
+    super(_PresubmitAddReviewers, self).__init__('')
+    self.reviewers = reviewers
+
+  def handle(self, output):
+    output.reviewers.extend(self.reviewers)
+
+
+# Top level object so multiprocessing can pickle
+# Public access through OutputApi object.
+class _PresubmitError(_PresubmitResult):
+  """A hard presubmit error."""
+  fatal = True
+
+
+# Top level object so multiprocessing can pickle
+# Public access through OutputApi object.
+class _PresubmitPromptWarning(_PresubmitResult):
+  """An warning that prompts the user if they want to continue."""
+  should_prompt = True
+
+
+# Top level object so multiprocessing can pickle
+# Public access through OutputApi object.
+class _PresubmitNotifyResult(_PresubmitResult):
+  """Just print something to the screen -- but it's not even a warning."""
+  pass
+
+
+# Top level object so multiprocessing can pickle
+# Public access through OutputApi object.
+class _MailTextResult(_PresubmitResult):
+  """A warning that should be included in the review request email."""
+  def __init__(self, *args, **kwargs):
+    super(_MailTextResult, self).__init__()
+    raise NotImplementedError()
+
+
 class OutputApi(object):
   """An instance of OutputApi gets passed to presubmit scripts so that they
   can output various types of results.
   """
+  PresubmitResult = _PresubmitResult
+  PresubmitAddReviewers = _PresubmitAddReviewers
+  PresubmitError = _PresubmitError
+  PresubmitPromptWarning = _PresubmitPromptWarning
+  PresubmitNotifyResult = _PresubmitNotifyResult
+  MailTextResult = _MailTextResult
+
   def __init__(self, is_committing):
     self.is_committing = is_committing
 
-  class PresubmitResult(object):
-    """Base class for result objects."""
-    fatal = False
-    should_prompt = False
-
-    def __init__(self, message, items=None, long_text=''):
-      """
-      message: A short one-line message to indicate errors.
-      items: A list of short strings to indicate where errors occurred.
-      long_text: multi-line text output, e.g. from another tool
-      """
-      self._message = message
-      self._items = []
-      if items:
-        self._items = items
-      self._long_text = long_text.rstrip()
-
-    def handle(self, output):
-      output.write(self._message)
-      output.write('\n')
-      for index, item in enumerate(self._items):
-        output.write('  ')
-        # Write separately in case it's unicode.
-        output.write(str(item))
-        if index < len(self._items) - 1:
-          output.write(' \\')
-        output.write('\n')
-      if self._long_text:
-        output.write('\n***************\n')
-        # Write separately in case it's unicode.
-        output.write(self._long_text)
-        output.write('\n***************\n')
-      if self.fatal:
-        output.fail()
-
-  class PresubmitAddReviewers(PresubmitResult):
-    """Add some suggested reviewers to the change."""
-    def __init__(self, reviewers):
-      super(OutputApi.PresubmitAddReviewers, self).__init__('')
-      self.reviewers = reviewers
-
-    def handle(self, output):
-      output.reviewers.extend(self.reviewers)
-
-  class PresubmitError(PresubmitResult):
-    """A hard presubmit error."""
-    fatal = True
-
-  class PresubmitPromptWarning(PresubmitResult):
-    """An warning that prompts the user if they want to continue."""
-    should_prompt = True
-
-  class PresubmitNotifyResult(PresubmitResult):
-    """Just print something to the screen -- but it's not even a warning."""
-    pass
-
   def PresubmitPromptOrNotify(self, *args, **kwargs):
     """Warn the user when uploading, but only notify if committing."""
     if self.is_committing:
       return self.PresubmitNotifyResult(*args, **kwargs)
     return self.PresubmitPromptWarning(*args, **kwargs)
 
-  class MailTextResult(PresubmitResult):
-    """A warning that should be included in the review request email."""
-    def __init__(self, *args, **kwargs):
-      super(OutputApi.MailTextResult, self).__init__()
-      raise NotImplementedError()
-
 
 class InputApi(object):
   """An instance of this object is passed to presubmit scripts so they can
   know stuff about the change they're looking at.
   """
   # Method could be a function
   # pylint: disable=R0201
 
   # File extensions that are considered source files from a style guide
   # perspective. Don't modify this list from a presubmit script!
@@ skipping 85 matching lines @@
     self._current_presubmit_path = os.path.dirname(presubmit_path)
 
     # We carry the canned checks so presubmit scripts can easily use them.
     self.canned_checks = presubmit_canned_checks
 
     # TODO(dpranke): figure out a list of all approved owners for a repo
     # in order to be able to handle wildcard OWNERS files?
     self.owners_db = owners.Database(change.RepositoryRoot(),
         fopen=file, os_path=self.os_path, glob=self.glob)
     self.verbose = verbose
+    self.Command = CommandData
 
     # Replace <hash_map> and <hash_set> as headers that need to be included
     # with "base/hash_tables.h" instead.
     # Access to a protected member _XX of a client class
     # pylint: disable=W0212
     self.cpplint._re_pattern_templates = [
       (a, b, 'base/hash_tables.h')
         if header in ('<hash_map>', '<hash_set>') else (a, b, header)
       for (a, b, header) in cpplint._re_pattern_templates
     ]
@@ skipping 133 matching lines @@
       file_item = file_item.AbsoluteLocalPath()
     if not file_item.startswith(self.change.RepositoryRoot()):
       raise IOError('Access outside the repository root is denied.')
     return gclient_utils.FileRead(file_item, mode)
 
   @property
   def tbr(self):
     """Returns if a change is TBR'ed."""
     return 'TBR' in self.change.tags
 
+  @staticmethod
+  def RunTests(tests_mix, parallel=True):
+    tests = []
+    msgs = []
+    for t in tests_mix:
+      if isinstance(t, OutputApi.PresubmitResult):
+        msgs.append(t)
+      else:
+        assert issubclass(t.message, _PresubmitResult)
+        tests.append(t)
+    if parallel:
+      pool = multiprocessing.Pool()
+      # async recipe works around multiprocessing bug handling Ctrl-C
+      msgs.extend(pool.map_async(CallCommand, tests).get(99999))
+      pool.close()
+      pool.join()
+    else:
+      msgs.extend(map(CallCommand, tests))
+    return [m for m in msgs if m]
+
 
 class AffectedFile(object):
   """Representation of a file in a change."""
   # Method could be a function
   # pylint: disable=R0201
   def __init__(self, path, action, repository_root):
     self._path = path
     self._action = action
     self._local_root = repository_root
     self._is_directory = None
@@ skipping 781 matching lines @@
     for method_name in method_names:
       if not hasattr(presubmit_canned_checks, method_name):
         raise NonexistantCannedCheckFilter(method_name)
       filtered[method_name] = getattr(presubmit_canned_checks, method_name)
       setattr(presubmit_canned_checks, method_name, lambda *_a, **_kw: [])
     yield
   finally:
     for name, method in filtered.iteritems():
       setattr(presubmit_canned_checks, name, method)
 
+def CallCommand(cmd_data):
+  # multiprocessing needs a top level function with a single argument.
+  cmd_data.kwargs['stdout'] = subprocess.PIPE
+  cmd_data.kwargs['stderr'] = subprocess.STDOUT
+  try:
+    (out, _), code = subprocess.communicate(cmd_data.cmd, **cmd_data.kwargs)
+    if code != 0:
+      return cmd_data.message('%s failed\n%s' % (cmd_data.name, out))
+  except OSError as e:
+    return cmd_data.message(
+        '%s exec failure\n   %s\n%s' % (cmd_data.name, e, out))
+
 
 def Main(argv):
   parser = optparse.OptionParser(usage="%prog [options] <files...>",
                                  version="%prog " + str(__version__))
   parser.add_option("-c", "--commit", action="store_true", default=False,
                    help="Use commit instead of upload checks")
   parser.add_option("-u", "--upload", action="store_false", dest='commit',
                    help="Use upload instead of commit checks")
   parser.add_option("-r", "--recursive", action="store_true",
                    help="Act recursively")
@@ skipping 60 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))