From 3ee6e2a34029bc5e68a5f733749e2732c2bf7940 Mon Sep 17 00:00:00 2001
From: google-java-format Team <java-team-github-bot@google.com>
Date: Thu, 28 Mar 2024 17:07:54 -0700
Subject: [PATCH] Java-format: format multiples files in parallel to improve
 speed

PiperOrigin-RevId: 620100015
---
 scripts/google-java-format-diff.py | 82 ++++++++++++++++++------------
 1 file changed, 49 insertions(+), 33 deletions(-)

diff --git a/scripts/google-java-format-diff.py b/scripts/google-java-format-diff.py
index 2c75edac7..7f52ed1ec 100755
--- a/scripts/google-java-format-diff.py
+++ b/scripts/google-java-format-diff.py
@@ -33,8 +33,34 @@
 import subprocess
 import io
 import sys
+from concurrent.futures import ThreadPoolExecutor,wait,FIRST_EXCEPTION
 from shutil import which
 
+def _apply_format(filename, lines, base_command, args):
+  """Apply format on filename."""
+  if args.i and args.verbose:
+    print('Formatting', filename)
+
+  command = base_command[:]
+  command.extend(lines)
+  command.append(filename)
+  p = subprocess.Popen(command, stdout=subprocess.PIPE,
+                       stderr=None, stdin=subprocess.PIPE)
+  stdout, _ = p.communicate()
+  if p.returncode != 0:
+    sys.exit(p.returncode)
+
+  if not args.i:
+    with open(filename) as f:
+      code = f.readlines()
+    formatted_code = io.StringIO(stdout.decode('utf-8')).readlines()
+    diff = difflib.unified_diff(code, formatted_code,
+                                filename, filename,
+                                '(before formatting)', '(after formatting)')
+    diff_string = ''.join(diff)
+    if len(diff_string) > 0:
+      sys.stdout.write(diff_string)
+
 def main():
   parser = argparse.ArgumentParser(description=
                                    'Reformat changed lines in diff. Without -i '
@@ -108,39 +134,29 @@ def main():
     binary = which('google-java-format') or '/usr/bin/google-java-format'
     base_command = [binary]
 
-  # Reformat files containing changes in place.
-  for filename, lines in lines_by_file.items():
-    if args.i and args.verbose:
-      print('Formatting', filename)
-    command = base_command[:]
-    if args.i:
-      command.append('-i')
-    if args.aosp:
-      command.append('--aosp')
-    if args.skip_sorting_imports:
-      command.append('--skip-sorting-imports')
-    if args.skip_removing_unused_imports:
-      command.append('--skip-removing-unused-imports')
-    if args.skip_javadoc_formatting:
-      command.append('--skip-javadoc-formatting')
-    command.extend(lines)
-    command.append(filename)
-    p = subprocess.Popen(command, stdout=subprocess.PIPE,
-                         stderr=None, stdin=subprocess.PIPE)
-    stdout, stderr = p.communicate()
-    if p.returncode != 0:
-      sys.exit(p.returncode);
-
-    if not args.i:
-      with open(filename) as f:
-        code = f.readlines()
-      formatted_code = io.StringIO(stdout.decode('utf-8')).readlines()
-      diff = difflib.unified_diff(code, formatted_code,
-                                  filename, filename,
-                                  '(before formatting)', '(after formatting)')
-      diff_string = ''.join(diff)
-      if len(diff_string) > 0:
-        sys.stdout.write(diff_string)
+  if args.i:
+    base_command.append('-i')
+  if args.aosp:
+    base_command.append('--aosp')
+  if args.skip_sorting_imports:
+    base_command.append('--skip-sorting-imports')
+  if args.skip_removing_unused_imports:
+    base_command.append('--skip-removing-unused-imports')
+  if args.skip_javadoc_formatting:
+    base_command.append('--skip-javadoc-formatting')
+
+  with ThreadPoolExecutor() as executor:
+    format_futures = []
+    for filename, lines in lines_by_file.items():
+      format_futures.append(
+          executor.submit(_apply_format, filename, lines, base_command, args)
+      )
+
+    done, _ = wait(format_futures, return_when=FIRST_EXCEPTION)
+    for future in done:
+      if exception := future.exception():
+        executor.shutdown(wait=True, cancel_futures=True)
+        sys.exit(exception.args[0])
 
 if __name__ == '__main__':
   main()