prepare_upgrade.py

#!/usr/bin/env python3

#
# Copyright 2023, The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Script to prepare an update to a new version of ktfmt."""

import subprocess
import os
import sys
import re
import shutil
import argparse
import textwrap

tmp_dir = "/tmp/ktfmt"
zip_path = os.path.join(tmp_dir, "common.zip")
jar_path = os.path.join(tmp_dir, "framework/ktfmt.jar")
copy_path = os.path.join(tmp_dir, "copy.jar")


def main():
  parser = argparse.ArgumentParser(
      description="Prepare a repository for the upgrade of ktfmt to a new version."
  )
  parser.add_argument(
      "--build_id",
      required=True,
      help="The build ID of aosp-build-tools-release with the new version of ktfmt"
  )
  parser.add_argument(
      "--bug_id",
      required=True,
      help="The bug ID associated to each CL generated by this tool")
  parser.add_argument(
      "--repo",
      required=True,
      help="The relative path of the repository to upgrade, e.g. 'frameworks/base/'"
  )
  args = parser.parse_args()

  build_id = args.build_id
  bug_id = args.bug_id
  repo_relative_path = args.repo

  build_top = os.environ["ANDROID_BUILD_TOP"]
  repo_absolute_path = os.path.join(build_top, repo_relative_path)

  print("Preparing upgrade of ktfmt from build", build_id)
  os.chdir(repo_absolute_path)
  check_workspace_clean()
  check_branches()

  print("Downloading ktfmt.jar from aosp-build-tools-release")
  download_jar(build_id)

  print(f"Creating local branch ktfmt_update1")
  run_cmd(["repo", "start", "ktfmt_update1"])

  includes_file = find_includes_file(repo_relative_path)
  if includes_file:
    update_includes_file(build_top, includes_file, bug_id)
  else:
    print("No includes file found, skipping first CL")

  print(f"Creating local branch ktfmt_update2")
  run_cmd(["repo", "start", "--head", "ktfmt_update2"])
  format_files(build_top, includes_file, repo_absolute_path, bug_id)

  print("Done. You can now submit the generated CL(s), if any.")


def run_cmd(cmd):
  result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  if result.returncode != 0:
    print("Error running command: {}".format(" ".join(cmd)))
    print("Output: {}".format(result.stderr.decode()))
    sys.exit(1)
  return result.stdout.decode("utf-8")


def is_workspace_clean():
  return run_cmd(["git", "status", "--porcelain"]) == ""


def check_workspace_clean():
  if not is_workspace_clean():
    print(
        "The current repository contains uncommitted changes, please run this script in a clean workspace"
    )
    sys.exit(1)


def check_branches():
  result = run_cmd(["git", "branch"])
  if "ktfmt_update1" in result or "ktfmt_update2" in result:
    print(
        "Branches ktfmt_update1 or ktfmt_update2 already exist, you should delete them before running this script"
    )
    sys.exit(1)


def download_jar(build_id):
  cmd = [
      "/google/data/ro/projects/android/fetch_artifact", "--branch",
      "aosp-build-tools-release", "--bid", build_id, "--target", "linux",
      "build-common-prebuilts.zip", zip_path
  ]
  run_cmd(cmd)
  cmd = ["unzip", "-q", "-o", "-d", tmp_dir, zip_path]
  run_cmd(cmd)

  if not os.path.isfile(jar_path):
    print("Error: {} is not readable".format(jar_path))
    sys.exit(1)


def find_includes_file(repo_relative_path):
  with open("PREUPLOAD.cfg") as f:
    includes_line = [line for line in f if "ktfmt.py" in line][0].split(" ")
    if "-i" not in includes_line:
      return None

    index = includes_line.index("-i") + 1
    includes_file = includes_line[index][len("${REPO_ROOT}/") +
                                         len(repo_relative_path):]
    if not os.path.isfile(includes_file):
      print("Error: {} does not exist or is not a file".format(includes_file))
      sys.exit(1)
    return includes_file


def get_included_folders(includes_file):
  with open(includes_file) as f:
    return [line[1:] for line in f.read().splitlines() if line.startswith("+")]


def update_includes_file(build_top, includes_file, bug_id):
  included_folders = get_included_folders(includes_file)
  cmd = [
      f"{build_top}/external/ktfmt/generate_includes_file.py",
      f"--output={includes_file}"
  ] + included_folders
  print(f"Updating {includes_file} with the command: {cmd}")
  run_cmd(cmd)

  if is_workspace_clean():
    print(f"No change were made to {includes_file}, skipping first CL")
  else:
    print(f"Creating first CL with update of {includes_file}")
    create_first_cl(bug_id)


def create_first_cl(bug_id):
  sha1sum = get_sha1sum(jar_path)
  change_id = f"I{sha1sum}"
  command = " ".join(sys.argv)
  cl_message = textwrap.dedent(f"""
  Regenerate include file for ktfmt upgrade

  This CL was generated automatically from the following command:

  $ {command}

  This CL regenerates the inclusion file with the current version of ktfmt
  so that it is up-to-date with files currently formatted or ignored by
  ktfmt.

  Bug: {bug_id}
  Test: Presubmits
  Change-Id: {change_id}
  Merged-In: {change_id}
  """)

  run_cmd(["git", "add", "--all"])
  run_cmd(["git", "commit", "-m", cl_message])


def get_sha1sum(file):
  output = run_cmd(["sha1sum", file])
  regex = re.compile(r"[a-f0-9]{40}")
  match = regex.search(output)
  if not match:
    print(f"sha1sum not found in output: {output}")
    sys.exit(1)
  return match.group()


def format_files(build_top, includes_file, repo_absolute_path, bug_id):
  if (includes_file):
    included_folders = get_included_folders(includes_file)
    cmd = [
        f"{build_top}/external/ktfmt/ktfmt.py", "-i", includes_file, "--jar",
        jar_path
    ] + included_folders
  else:
    cmd = [
        f"{build_top}/external/ktfmt/ktfmt.py", "--jar", jar_path,
        repo_absolute_path
    ]

  print(
      f"Formatting the files that are already formatted with the command: {cmd}"
  )
  run_cmd(cmd)

  if is_workspace_clean():
    print("All files were already properly formatted, skipping second CL")
  else:
    print("Creating second CL that formats all files")
    create_second_cl(bug_id)


def create_second_cl(bug_id):
  # Append 'ktfmt_update' at the end of a copy of the jar file to get
  # a different sha1sum.
  shutil.copyfile(jar_path, copy_path)
  with open(copy_path, "a") as file_object:
    file_object.write("ktfmt_update")

  sha1sum = get_sha1sum(copy_path)
  change_id = f"I{sha1sum}"
  command = " ".join(sys.argv)
  cl_message = textwrap.dedent(f"""
  Format files with the upcoming version of ktfmt

  This CL was generated automatically from the following command:

  $ {command}

  This CL formats all files already correctly formatted with the upcoming
  version of ktfmt.

  Bug: {bug_id}
  Test: Presubmits
  Change-Id: {change_id}
  Merged-In: {change_id}
  """)

  run_cmd(["git", "add", "--all"])
  run_cmd(["git", "commit", "-m", cl_message])


if __name__ == "__main__":
  main()