commit_to_hg.py

#!/usr/bin/env python

import argparse
from email.policy import default
import os
import pathlib
import re
import shutil
import subprocess
import sys

from pathlib import Path

from time import time
from typing import List, Tuple

# This script allows you to forget about HG as much as is possible

# Things to know
# If an error is encountered in any commit, the script will exit
# 'arc f' changes are not copied back to git so stacked commits will eventually get out of sync and fail

HOME = pathlib.Path.home()
UNIXNAME = os.environ['USER']
HG_IOS_SDK_DIR = f"{HOME}/fbsource/fbobjc/ios-sdk"

# DEFINE THE LOCATION OF THE GIT DIR
GIT_IOS_SDK_DIR = f"{HOME}/github.com/facebook/facebook-ios-sdk"

if UNIXNAME == 'jawwad':
    TEST_IPHONE='iPhone 12'
else:
    TEST_IPHONE='iPhone 13 Pro Max'

HG_COMMIT_MESSAGE_FILE = "/tmp/hg_commit_message.txt"
PATCH_FILE = "git_diff.patch"

def parse_arguments():
    # https://docs.python.org/3/library/argparse.html
    parser = argparse.ArgumentParser()
    # Important Options
    parser.add_argument("--sha", default='head', help="Commit the SHA to HG")
    parser.add_argument("--revert", "-r", action="store_true", help="Revert any uncommitted changes in HG")
    parser.add_argument("--skip-lint", action="store_true", help="Skip Linting")
    parser.add_argument("--skip-tests", action="store_true", help="Skip Tests")
    parser.add_argument("--skip-submit", action="store_true", help="Skip Submitting as a diff")
    parser.add_argument("--skip-checkout-master", action="store_true", help="Skip checking out latest branch")
    parser.add_argument("--delete-branches", action="store_true", help="Delete branches for committed diffs")

    # Experimental Options
    parser.add_argument("--update-gist", action="store_true", help="Skip Linting")
    parser.add_argument("--self-update", action="store_true", help="Skip Linting")
    parser.add_argument("--copy-only", action="store_true", help="Only copy changes and deletions from Git to HG")
    parser.add_argument("--import-diff", help="Experimental: The diff to import")
    parser.add_argument("--reverse-copy", action="store_true", help="Copy changes and deletions from HG to Git")

    group = parser.add_mutually_exclusive_group(required=False)
    group.add_argument("--amend", action="store_true")
    group.add_argument("--stack", action="store_true")

    return parser.parse_args()


def create_and_apply_patch(sha=None):
    if sha:
        patch_command = f"git format-patch -1 {sha} --stdout"
    else:
        patch_command = f"git diff --cached"

    # Note: This fails if there isn't anything staged
    run(f"{patch_command} > {HG_IOS_SDK_DIR}/{PATCH_FILE}")  # make a patch from staged files
    run(f"hg import {PATCH_FILE} --no-commit --prefix .")
    run(f"rm {PATCH_FILE}")


def notify_about_commit_message(sha):
    if sha:
        message = commit_message_from_sha(sha)
        fields = ["Summary", "Test Plan", "Reviewers", "Subscribers", "Tasks", "Tags"]
        for field in fields:
            message = message.replace(f"{field}:\n\n", f"{field}:\n")

        write_text_to_file(message, HG_COMMIT_MESSAGE_FILE)

    print_yellow(f"\nTHE COMMIT WILL BE MADE WITH THE FOLLOWING MESSAGE. (Control-C if this is incorrect)")

    template_text = read_text_from_file(HG_COMMIT_MESSAGE_FILE)

    # Simplify output by stripping empty fields
    # fields = ["Summary", "Test Plan", "Reviewers", "Subscribers", "Tasks", "Tags"]
    # for field in fields:
        # template_text = template_text.replace(f"{field}:\n\n", f"{field}:\n")

    divider = "-" * 80
    print_yellow(divider)
    print(template_text)
    print_yellow(divider)

def import_diff_from_hg(diff: str) -> None:
    run(f"git checkout main")
    run(f"hg checkout {diff}")
    run(f"hg diff  --change . > {GIT_IOS_SDK_DIR}/{PATCH_FILE}")
    os.chdir(GIT_IOS_SDK_DIR)
    run(f"perl -pi -e s'/fbobjc\/ios-sdk\///' {PATCH_FILE}")  # strip fbobjc/ios-sdk/
    output = run(f"git apply --reject {PATCH_FILE}")  # reject skips things that can't be applied
    run(f"rm {PATCH_FILE}")
    print(f"Output: {output}")
    run(f"git add .")
    run(f"git checkout -b 'diffs/{diff}'")
    run(f"git commit -m 'Imported Diff: {diff}'")

def hg_revert_and_purge():
    run("hg revert --all")  # revert tracked changes
    run("hg purge --files")  # delete untracked files

def update_gist():
    gist_url = "https://gist.github.com/02dba3eb06a2429089ef6db82b6756a6"
    script_path = Path(__file__).absolute()
    run(f"gist --filename commit_to_hg.py --update {gist_url} < '{script_path}'")

def self_update():
    # https://gist.github.com/cubedtear/54434fc66439fc4e04e28bd658189701
    gist_url = "https://gist.github.com/02dba3eb06a2429089ef6db82b6756a6"
    script_dir = Path(__file__).parent.absolute()
    # run(f"curl -O https://gist.githubusercontent.com/jawwad/02dba3eb06a2429089ef6db82b6756a6/raw/commit_to_hg.py")
    run(f"curl -o /tmp/script.py https://gist.githubusercontent.com/jawwad/02dba3eb06a2429089ef6db82b6756a6/raw/commit_to_hg.py")
    run(f"chmod 755 /tmp/script.py")
    script_path = Path(__file__).absolute()
    print("Run the following to update this file:")
    print(f"mv /tmp/script.py '{script_path}'")

def main():
    args = parse_arguments()

    if args.import_diff:
        import_diff_from_hg(args.import_diff)
        sys.exit(0)
    elif args.update_gist:
        update_gist()
        sys.exit(0)
    elif args.self_update:
        self_update()
        sys.exit(0)
    elif args.delete_branches:
        run(f"git checkout main")
        delete_branches_for_committed_diffs()
        sys.exit(0)
    elif args.copy_only:
        copy_only(args)
        sys.exit(0)
    elif args.reverse_copy:
        reverse_copy_only()
        sys.exit(0)

    if args.amend:
        print_yellow(f"AMENDING EXISTING COMMIT")
    else:
        notify_about_commit_message(args.sha)

    hg_status_output = run("hg status --color=always")
    if hg_status_output:
        if args.revert:
            hg_revert_and_purge()
        else:
            print_red("UNCOMMITTED CHANGES EXIST IN HG. PLEASE COMMIT OR REVERT THESE CHANGES (hg revert --all && hg purge --files)")
            print(hg_status_output)
            revert_or_abort = input(f"Type 'r' to revert these changes or any other key to abort: ")
            if revert_or_abort == "r" or revert_or_abort == "R":
                hg_revert_and_purge()
            else:
                print_error_and_exit("Exiting")

    should_amend = False

    git_branch_name = current_git_branch_name()
    on_diffs_branch = re.search(r"diffs/(D\d+)", git_branch_name)
    if on_diffs_branch:
        diff_number = on_diffs_branch.group(1)
        run(f"hg checkout {diff_number}")

        print_yellow(f"ON DIFFS BRANCH FOR DIFF: {diff_number}")

        if args.amend:
            should_amend = True
        elif not args.stack:
            amend_or_commit = input(f"Type 'a' to amend to {diff_number} or 's' to create a new stacked commit: ")
            if amend_or_commit == "a":
                should_amend = True
            elif amend_or_commit != "s":
                print_error_and_exit("Invalid Input. Exiting")
    else:
        # If branch is not main, just notify but still proceed as normal
        # if git_branch_name != "main":
        #     print_yellow(f"Current git branch is: {git_branch_name}")

        if not args.skip_checkout_master:
            run("hg pull")
            run("hg checkout master")

    if args.sha:
        create_and_apply_patch(args.sha)
    else:
        create_and_apply_patch("head")

    try_to_hg_mv_any_moved_files()

    # Uncomment and edit to include a custom manual rename
    # run("hg mv --after ObjectiveCFile.h SwiftFile.swift")

    arc_f_output = run("arc f").strip()
    if arc_f_output and arc_f_output != "ok No lint issues.":
        print_yellow(f"arc_f_output: {arc_f_output}|")
        sys.exit(0)

    if not args.skip_lint:
        arc_lint_output = run("arc lint --apply-patches").strip()
        if arc_lint_output != "ok No lint issues.":
            print_yellow(f"arc_lint_output: {arc_lint_output}|")
            sys.exit(0)

    run("./generate-projects.sh --skip-closing-xcode")

    if not args.skip_tests:
        run(f"xcodebuild clean test -workspace FacebookSDK.xcworkspace -scheme BuildAllKits-Dynamic -destination 'platform=iOS Simulator,name={TEST_IPHONE}'")

    if should_amend:
        run("hg amend")
    else:
        run(f"hg commit -l {HG_COMMIT_MESSAGE_FILE}")

    if args.skip_submit:
        return

    jf_submit_output = run("jf submit --draft")

    diff_match = re.search(r"https://www.internalfb.com/diff/(D\d+)", jf_submit_output)
    if diff_match:
        diff_number = diff_match.group(1)
        if should_amend:
            run(f"git commit --amend --no-edit")
        else:
            run(f"git checkout -b diffs/{diff_number}")
        run(f"open https://www.internalfb.com/diff/{diff_number}")
    else:
        print(f"Couldn't find diff number from hg commit output: {jf_submit_output}")

    delete_branches_for_committed_diffs()


def try_to_hg_mv_any_moved_files():
    hg_status_output = run("hg status")
    added_file_map = {}
    deleted_file_map = {}
    for line in hg_status_output.splitlines():
        (status, path) = line.split()
        filename = os.path.basename(path)
        if status == "A":
            added_file_map[filename] = path
        elif status == "R":
            deleted_file_map[filename] = path

    for filename, swift_path in added_file_map.items():
        filename_without_extension = filename.replace(".swift", "").replace("_", "")

        possibly_deleted_files = [
            f"FBSDK{filename_without_extension}.m", # m must be first
            f"FBSDK{filename_without_extension}.h",
            f"_FBSDK{filename_without_extension}.m", # m must be first
            f"_FBSDK{filename_without_extension}.h",
            f"FBSDK{filename_without_extension}.swift", # prefix removed
            f"_{filename_without_extension}.swift", # _ removed
        ]

        deleted_file_path = None
        for file in possibly_deleted_files:
            if file in deleted_file_map:
                deleted_file_path = deleted_file_map[file]
                break

        if deleted_file_path:
            print_yellow("MOVE DETECTED:")
            run(f"hg mv --after {deleted_file_path} {swift_path}")


def commit_message_from_sha(sha): #title, summary, test_plan="sandcastle", tasks=""):
    # add newline before summary and test plan if on multiple lines

    title = get_git_log_subject(sha).strip()
    (summary, test_plan, tasks) = get_git_log_body(sha)

    return f"""\
{title}

Summary:
{summary}

Test Plan: {test_plan}

Reviewers: {get_reviewers()}

Tasks: {tasks}

Tags: accept2ship
"""

def get_git_log_subject(sha: str) -> str:
    return run(f"git log -n1 --pretty=format:%s {sha}")

def get_git_log_body(sha: str) -> Tuple[str, str, str]:
    body = run(f"git log -n1 --pretty=format:%b {sha}").strip()

    task = ""
    tasks_regex = r"Tasks?: (\w+)"
    task_match = re.search(tasks_regex, body, flags=re.MULTILINE)
    if task_match:
        task = task_match.group(1)
        body = re.sub(tasks_regex, r"", body)
        body = body.strip()

    if "Test Plan:" in body:
        (summary, test_plan) = body.split("Test Plan:")
    else:
        summary, test_plan = body, "sandcastle"

    return (summary.strip(), test_plan, task)

def get_reviewers() -> str:
    reviewers = ["jawwad", "samodom", "joesusnick"]
    reviewers.remove(UNIXNAME)
    reviewers.remove("joesusnick")
    return ", ".join(reviewers)

def delete_branches_for_committed_diffs():
    published_diffs = get_published_diffs()

    print_cyan("Checking status for published diffs")
    for diff in published_diffs:
        status = get_diff_status(diff)

        if status in ("Committed", "Abandoned", "Reverted"):
            print_yellow(f"DELETING BRANCH for {status} diff: {diff}")
            print(run(f"git branch -D diffs/{diff}"))
        else:
            print_yellow(f"DIFF: {diff}, STATUS: {status}")


def run(command, check=True):
    if command.startswith(("hg", "arc")):
        os.chdir(HG_IOS_SDK_DIR)
    elif command.startswith("git"):
        os.chdir(GIT_IOS_SDK_DIR)

    print_cyan(f"RUNNING: {command}", end="", flush=True)  # without flush it waits to print after the command

    try:
        start = time()
        completed_process = subprocess.run(command, shell=True, check=check, capture_output=True)
        end = time()
        print(f" (completed in {end-start:.2f} sec)")

        bytes_str = completed_process.stdout
        return str(bytes_str, "utf-8")

    except subprocess.CalledProcessError as e:
        print(f"{type(e).__name__}: {e}")
        print(f"ERROR Running Command: {command}")
        print(f"STDOUT: {e.stdout.decode()}")
        print(f"STDERR: {e.stderr.decode()}")
        print(f"RETURNCODE: {e.returncode}")
        print(f"CMD: {e.cmd}")
        print_error_and_exit("Exiting due to error")


# Helper methods

def current_git_branch_name():
    return run("git branch --show-current").rstrip()


def read_text_from_file(file: str) -> str:
    with open(file, "r") as f:
        text = f.read()
    return text

def write_text_to_file(text: str, file: str) -> None:
    with open(file, "w") as f:
        f.write(text)


def get_diff_status(diff: str) -> str:
    return run(f"hg log -T'{{phabstatus}}' -r {diff}")


def get_published_diffs() -> List[str]:
    """Returns diff numbers from branches named diffs/DXXXXXX"""
    return run("git --no-pager branch --list 'diffs/*' | cut -d '/' -f 2").splitlines()


# Print helpers
def print_cyan(message, **kwargs):
    _print_color(Colors.CYAN, message, **kwargs)


def print_red(message, **kwargs):
    _print_color(Colors.RED, message, **kwargs)


def print_green(message, **kwargs):
    _print_color(Colors.GREEN, message, **kwargs)


def print_yellow(message, **kwargs):
    _print_color(Colors.YELLOW, message, **kwargs)


def _print_color(color, message, **kwargs):
    print(color + message + Colors.RESET, **kwargs)


def print_error_and_exit(*args, **kwargs):
    print_red(*args, file=sys.stderr, **kwargs)
    sys.exit(1)


def copy_only(args):
    """Commit changed files from Git repo to HG repo"""
    uncommitted_changes = run("hg status")
    if uncommitted_changes and not args.revert:
        print("Uncommitted changes exist in HG. Please commit or revert these changes by running: hg revert --all && hg purge --files")
        print(uncommitted_changes)
        sys.exit(1)

    if args.revert:
        hg_revert_and_purge()

    output = run("git status -s --porcelain")

    os.chdir(GIT_IOS_SDK_DIR)
    for line in output.splitlines():
        (status, path) = line.strip().split()
        destpath = f"{HG_IOS_SDK_DIR}/{path}"

        if status in ("A", "AM", "M", "??", "MM"):
            shutil.copy(path, destpath)
        elif status == "D":
            if os.path.exists(destpath):
                os.remove(destpath)
        else:
            print(f"Unknown status: {status} for file {path}")

    lint_output = run("arc f && arc lint")
    print(lint_output)


def reverse_copy_only():
    """Commit changed files from HG repo to Git repo"""
    os.chdir(HG_IOS_SDK_DIR)
    # output = run("hg status --rev bottom::top")
    # output = run("hg status --rev 'first(. % master)~1'") # This didn't work for uncommitted files
    output = run("hg status")

    for line in output.splitlines():
        (status, path) = line.strip().split()
        destpath = f"{GIT_IOS_SDK_DIR}/{path}"

        if status in ("A", "M", "??", "MM", "?"):
            shutil.copy(path, destpath)
        elif status == "D":
            if os.path.exists(destpath):
                os.remove(destpath)
        else:
            print(f"Unknown status: {status} for file {path}")


# ANSI color definitions
class Colors:
    RESET = "\033[0m"
    BLACK = "\033[30m"
    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    WHITE = "\033[37m"


if __name__ == "__main__":
    sys.exit(main())