tools/recontextualize_patches.py

#!/usr/bin/env python3
#
# Copyright (C) 2024 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.

"""Apply and re-generate our patches to decrease context drift"""

import argparse
from pathlib import Path
import shutil
import sys
from tempfile import TemporaryDirectory

import context

from android_rust.paths import PATCHES_PATH, RUST_SOURCE_PATH
from android_rust import source_manager
from android_rust.utils import GitRepo

#
# Program logic
#


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser("Apply and re-generate our patches to decrease context drift")

    mode_group = parser.add_mutually_exclusive_group()
    mode_group.add_argument("--apply-only", action="store_true", help="Stop after applying patches")
    mode_group.add_argument(
        "--resume",
        action="store_true",
        help="Resume a recontextualize operation after patch errors have been manually resolved")

    return parser.parse_args()


def main() -> None:
    args = parse_args()

    patch_list = sorted(PATCHES_PATH.glob("**/rustc-*"), key=lambda p: p.name)
    source_manager.ensure_unique_patch_numbers(patch_list)

    git_out = GitRepo(RUST_SOURCE_PATH)

    # Check to make sure that our rustc repository has no stray changes
    status = git_out.status()
    if status:
        sys.exit(f"Modified files in {RUST_SOURCE_PATH}\n{status}")

    if not args.resume:
        error_msg = (
            "Please manually resolve any patch applications in `toolchain/rustc` using Git commands.  " +
            "Once that is done the recontextualization process can completed by re-invoking this script with the `--resume` flag."
        )
        git_out.am(patch_list, error_msg=error_msg)

    if args.apply_only:
        return

    with TemporaryDirectory() as temp_dir:
        temp_dir_path = Path(temp_dir)

        git_out.format_patches(f"HEAD~{len(patch_list)}", outdir=temp_dir_path)

        for orig, new in zip(patch_list, sorted(temp_dir_path.iterdir())):
            shutil.move(new, orig)

    git_out.reset_hard(len(patch_list))


if __name__ == "__main__":
    try:
        main()
        sys.exit(0)
    except:
        sys.exit(1)