push-rewrite

#! /usr/bin/env python3

# push-rewrite -- Force push to a PR after rewriting history, with benefits.
#
# This tool force pushes your local changes to origin but only if that
# doesn't change any files.  If that is the case, the tool will also
# copy the test results over.
#
# The idea is that you use this tool after squashing fixups in a pull
# request or after other similar last minute rewriting activity before
# merging a pull request.  Then you can merge the PR using the GitHub
# UI and get a nice "merged" label for it, and the tests will still be
# green.

import argparse
import subprocess
import sys
import time

from lib.aio.jsonutil import JsonObject, get_dict, get_str
from task import github, labels_of_pull


def execute(*args: str) -> str:
    try:
        sys.stderr.write("+ " + " ".join(args) + "\n")
        output = subprocess.check_output(args, text=True)
    except subprocess.CalledProcessError as ex:
        sys.exit(ex.returncode)
    return output


def git(*args: str) -> str:
    return execute("git", *args).strip()


def find_pr_with_sha(api: github.GitHub, sha: str) -> JsonObject:
    pulls = api.pulls()
    for pull in pulls:
        if get_str(get_dict(pull, "head"), "sha") == sha:
            return pull
    sys.stderr.write(f"Could not find pull with revision {sha}.\n")
    sys.exit(1)


def main() -> None:
    parser = argparse.ArgumentParser(description='Force push after a rewrite')
    parser.add_argument('--repo', help="The GitHub repository to work with", default=None)
    opts = parser.parse_args()

    local = git('rev-parse', 'HEAD')
    remote = git('rev-parse', 'HEAD@{push}')

    if local == remote:
        sys.exit('Nothing to push')

    if git("diff", "--stat", local, remote) != "":
        sys.exit('You have local changes, aborting.')

    api = github.GitHub(repo=opts.repo)
    old_statuses = api.statuses(remote)

    pull = find_pr_with_sha(api, remote)
    labels = labels_of_pull(pull)
    tests_disabled = "no-test" in labels or "[no-test]" in get_str(pull, "title")
    # needs no-test label to prevent tests bring triggered
    # we cannot set the label for ourselves without `repo` permission
    if not tests_disabled:
        sys.exit("Please set the 'no-test' label on the PR before trying this")

    git("push", "--force-with-lease")

    for n in range(100, 0, -1):
        try:
            if api.get("commits/%s" % local):
                break
        except RuntimeError as e:
            if "Unprocessable Entity" not in str(e) or n <= 1:
                raise
            print("(new commits not yet in the GiHub API...please stand by. *beep*)")
        time.sleep(1)

    for key in old_statuses:
        if old_statuses[key]["state"] != "pending":
            print("Copying results for %s" % old_statuses[key]["context"])
            api.post("statuses/" + local, old_statuses[key])

    # remove no-test label
    if not tests_disabled:
        api.delete("issues/%i/labels/no-test" % pull["number"])


if __name__ == '__main__':
    main()