Refactor the way that SCM changes are tracked and aggregated in rebaseline.py.

third_party/WebKit/Tools/Scripts/webkitpy/tool/commands/rebaseline.py

 # Copyright (c) 2010 Google Inc. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are
 # met:
 #
 #     * Redistributions of source code must retain the above copyright
 # notice, this list of conditions and the following disclaimer.
 #     * Redistributions in binary form must reproduce the above
 # copyright notice, this list of conditions and the following disclaimer
@@ skipping 45 matching lines @@
     platform_options = factory.platform_options(use_globs=True)
 
     results_directory_option = optparse.make_option("--results-directory", help="Local results directory to use.")
 
     suffixes_option = optparse.make_option("--suffixes", default=','.join(BASELINE_SUFFIX_LIST), action="store",
                                            help="Comma-separated-list of file types to rebaseline.")
 
     def __init__(self, options=None):
         super(AbstractRebaseliningCommand, self).__init__(options=options)
         self._baseline_suffix_list = BASELINE_SUFFIX_LIST
-        self._scm_changes = {'add': [], 'delete': [], 'remove-lines': []}
+        self._scm_changes = ChangeSet()
         self._tool = None
 
-    def _add_to_scm_later(self, path):
-        self._scm_changes['add'].append(path)
+    def _print_scm_changes(self):
+        print(json.dumps(self._scm_changes.to_dict()))
 
-    def _delete_from_scm_later(self, path):
-        self._scm_changes['delete'].append(path)
 
-    def _print_scm_changes(self):
-        print(json.dumps(self._scm_changes))
+class ChangeSet(object):
+    """A record of added and deleted files and TestExpectation lines to remove.
+
+    This class groups two kinds of changes together:
+     (1) files that have been added or removed.
+     (2) lines to remove from TestExpectations.
+
+    The reason for keeping track of these changes together is that after rebaselining
+    many tests in parallel by subprocesses, we can update the git staged files list
+    and the TestExpectations all at once at the end.
+
+    Each subprocess can communicate some set of changes by outputting JSON lines,
+    and another process can aggregate the changes.
+
+    TODO(qyearsley): Refactor this so that:
+     - lines to remove are only tracked in one format.
+     - files to add and delete are always disjoint sets.
+    """
+    def __init__(self, files_to_add=None, files_to_delete=None, lines_to_remove=None):
+        self.files_to_add = files_to_add or []
+        self.files_to_delete = files_to_delete or []
+        self.lines_to_remove = lines_to_remove or {}
+
+    def add_file(self, path):
+        self.files_to_add.append(path)
+
+    def delete_file(self, path):
+        self.files_to_delete.append(path)
+
+    def remove_line(self, test, builder):
+        if test not in self.lines_to_remove:
+            self.lines_to_remove[test] = []
+        self.lines_to_remove[test].append(builder)
+
+    def to_dict(self):
+        remove_lines = []
+        for test in self.lines_to_remove:
+            for builder in self.lines_to_remove[test]:
+                remove_lines.append({'test': test, 'builder': builder})
+        return {
+            'add': list(self.files_to_add),
+            'delete': list(self.files_to_delete),
+            'remove-lines': remove_lines,
+        }
+
+    @staticmethod
+    def from_dict(change_dict):
+        files_to_add = set()
+        files_to_delete = set()
+        lines_to_remove = {}
+        if 'add' in change_dict:
+            files_to_add.update(change_dict['add'])
+        if 'delete' in change_dict:
+            files_to_delete.update(change_dict['delete'])
+        if 'remove-lines' in change_dict:
+            for line_to_remove in change_dict['remove-lines']:
+                test = line_to_remove['test']
+                builder = line_to_remove['builder']
+                if test not in lines_to_remove:
+                    lines_to_remove[test] = []
+                lines_to_remove[test].append(builder)
+        return ChangeSet(
+            files_to_add=list(files_to_add),
+            files_to_delete=list(files_to_delete),
+            lines_to_remove=lines_to_remove)
+
+    def update(self, other):
+        assert isinstance(other, ChangeSet)
+        assert type(other.lines_to_remove) is dict
+        self.files_to_add.extend(other.files_to_add)
+        self.files_to_delete.extend(other.files_to_delete)
+        for test in other.lines_to_remove:
+            if test not in self.lines_to_remove:
+                self.lines_to_remove[test] = []
+            self.lines_to_remove[test].extend(other.lines_to_remove[test])
 
 
 class BaseInternalRebaselineCommand(AbstractRebaseliningCommand):
     """Base class for rebaseline-related commands that are intended to be used by other commands."""
     # Not overriding execute() - pylint: disable=abstract-method
 
     def __init__(self):
         super(BaseInternalRebaselineCommand, self).__init__(options=[
             self.results_directory_option,
             self.suffixes_option,
@@ skipping 76 matching lines @@
             new_baselines.append(new_baseline)
 
         for i in range(len(old_baselines)):
             old_baseline = old_baselines[i]
             new_baseline = new_baselines[i]
 
             _log.debug("Copying baseline from %s to %s.", old_baseline, new_baseline)
             self._tool.filesystem.maybe_make_directory(self._tool.filesystem.dirname(new_baseline))
             self._tool.filesystem.copyfile(old_baseline, new_baseline)
             if not self._tool.scm().exists(new_baseline):
-                self._add_to_scm_later(new_baseline)
+                self._scm_changes.add_file(new_baseline)
 
     def execute(self, options, args, tool):
         self._tool = tool
         for suffix in options.suffixes.split(','):
             self._copy_existing_baseline(options.builder, options.test, suffix)
         self._print_scm_changes()
 
 
 class RebaselineTest(BaseInternalRebaselineCommand):
     name = "rebaseline-test-internal"
     help_text = "Rebaseline a single test from a buildbot. Only intended for use by other webkit-patch commands."
 
     def _save_baseline(self, data, target_baseline):
         if not data:
             _log.debug("No baseline data to save.")
             return
 
         filesystem = self._tool.filesystem
         filesystem.maybe_make_directory(filesystem.dirname(target_baseline))
         filesystem.write_binary_file(target_baseline, data)
         if not self._tool.scm().exists(target_baseline):
-            self._add_to_scm_later(target_baseline)
+            self._scm_changes.add_file(target_baseline)
 
     def _rebaseline_test(self, builder_name, test_name, suffix, results_url):
         baseline_directory = self._baseline_directory(builder_name)
 
         source_baseline = "%s/%s" % (results_url, self._file_name_for_actual_result(test_name, suffix))
         target_baseline = self._tool.filesystem.join(baseline_directory, self._file_name_for_expected_result(test_name, suffix))
 
         _log.debug("Retrieving source %s for target %s.", source_baseline, target_baseline)
         self._save_baseline(self._tool.web.get_binary(source_baseline, convert_404_to_None=True),
                             target_baseline)
 
     def _rebaseline_test_and_update_expectations(self, options):
         self._baseline_suffix_list = options.suffixes.split(',')
 
         port = self._tool.port_factory.get_from_builder_name(options.builder)
         if port.reference_files(options.test):
             if 'png' in self._baseline_suffix_list:
                 _log.warning("Cannot rebaseline image result for reftest: %s", options.test)
                 return
             assert self._baseline_suffix_list == ['txt']
 
         if options.results_directory:
             results_url = 'file://' + options.results_directory
         else:
             results_url = self._tool.buildbot.results_url(options.builder, build_number=options.build_number)
 
         for suffix in self._baseline_suffix_list:
             self._rebaseline_test(options.builder, options.test, suffix, results_url)
-        self._scm_changes['remove-lines'].append({'builder': options.builder, 'test': options.test})
+        self._scm_changes.remove_line(test=options.test, builder=options.builder)
 
     def execute(self, options, args, tool):
         self._tool = tool
         self._rebaseline_test_and_update_expectations(options)
         self._print_scm_changes()
 
 
 class AbstractParallelRebaselineCommand(AbstractRebaseliningCommand):
     """Base class for rebaseline commands that do some tasks in parallel."""
     # Not overriding execute() - pylint: disable=abstract-method
@@ skipping 91 matching lines @@
                     if options.verbose:
                         cmd_line.append('--verbose')
                     copy_baseline_commands.append(
                         tuple([[self._tool.executable, path_to_webkit_patch, 'copy-existing-baselines-internal'] + cmd_line, cwd]))
                     rebaseline_commands.append(
                         tuple([[self._tool.executable, path_to_webkit_patch, 'rebaseline-test-internal'] + cmd_line, cwd]))
         return copy_baseline_commands, rebaseline_commands, lines_to_remove
 
     @staticmethod
     def _serial_commands(command_results):
-        files_to_add = set()
-        files_to_delete = set()
-        lines_to_remove = {}
+        """Parses the JSON lines from command output to extract SCM changes.
+
+        TODO(qyearsley): Refactor and rename this function.
+        """
+        change_set = ChangeSet()
         for output in [result[1].split('\n') for result in command_results]:
             file_added = False
             for line in output:
                 try:
                     if line:
                         parsed_line = json.loads(line)
-                        if 'add' in parsed_line:
-                            files_to_add.update(parsed_line['add'])
-                        if 'delete' in parsed_line:
-                            files_to_delete.update(parsed_line['delete'])
-                        if 'remove-lines' in parsed_line:
-                            for line_to_remove in parsed_line['remove-lines']:
-                                test = line_to_remove['test']
-                                builder = line_to_remove['builder']
-                                if test not in lines_to_remove:
-                                    lines_to_remove[test] = []
-                                lines_to_remove[test].append(builder)
+                        change_set.update(ChangeSet.from_dict(parsed_line))
                         file_added = True
                 except ValueError:
                     _log.debug('"%s" is not a JSON object, ignoring', line)
 
             if not file_added:
                 _log.debug('Could not add file based off output "%s"', output)
 
-        return list(files_to_add), list(files_to_delete), lines_to_remove
+        return change_set
 
     def _optimize_baselines(self, test_prefix_list, verbose=False):
         optimize_commands = []
         for test in test_prefix_list:
             all_suffixes = set()
             builders_to_fetch_from = self._builders_to_fetch_from(self._builder_names(test_prefix_list[test]))
             for build in sorted(test_prefix_list[test]):
                 if build.builder_name not in builders_to_fetch_from:
                     break
                 all_suffixes.update(self._suffixes_for_actual_failures(test, build, test_prefix_list[test][build]))
@@ skipping 54 matching lines @@
             if fs.exists(smoke_test_filename) and test not in fs.read_text_file(smoke_test_filename):
                 return True
 
         return (SKIP in full_expectations.get_expectations(test) and
                 SKIP not in generic_expectations.get_expectations(test))
 
     def _run_in_parallel(self, commands, update_scm=True):
         if not commands:
             return {}
 
+        # TODO(qyearsley): Refactor this block.
         command_results = self._tool.executive.run_in_parallel(commands)
         log_output = '\n'.join(result[2] for result in command_results).replace('\n\n', '\n')
         for line in log_output.split('\n'):
             if line:
                 _log.error(line)
 
-        files_to_add, files_to_delete, lines_to_remove = self._serial_commands(command_results)
-        # TODO(qyearsley): Consider removing this if possible,
-        # or making it work and produce reasonable results for rebaseline-cl.
+        change_set = self._serial_commands(command_results)
+
+        # TODO(qyearsley): Instead of updating the SCM state here, aggregate changes
+        # and update once in _rebaseline. See http://crbug.com/639410.
         if update_scm:
-            if files_to_delete:
-                self._tool.scm().delete_list(files_to_delete)
-            if files_to_add:
-                self._tool.scm().add_list(files_to_add)
-        return lines_to_remove
+            if change_set.files_to_delete:
+                self._tool.scm().delete_list(change_set.files_to_delete)
+            if change_set.files_to_add:
+                self._tool.scm().add_list(change_set.files_to_add)
+
+        return change_set.lines_to_remove
 
     def _rebaseline(self, options, test_prefix_list, update_scm=True):
         """Downloads new baselines in parallel, then updates expectations files
         and optimizes baselines.
 
         Args:
             options: An object with the options passed to the current command.
             test_prefix_list: A map of test names to Build objects to file suffixes
                 for new baselines. For example:
                 {
@@ skipping 163 matching lines @@
             for test in args:
                 if test not in test_prefix_list:
                     test_prefix_list[test] = {}
                 build = Build(builder)
                 test_prefix_list[test][build] = suffixes_to_update
 
         if options.verbose:
             _log.debug("rebaseline-json: " + str(test_prefix_list))
 
         self._rebaseline(options, test_prefix_list)