Script to ghstack land

Author: kirklandsign

.github/scripts/propose_ghstack_orig_pr.py

@@ -0,0 +1,101 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import os
+import re
+
+from typing import List
+
+# Provided by the PyGithub pip package.
+from github import Auth, Github
+from github.Repository import Repository
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument(
+        "--repo",
+        type=str,
+        help='The github repo to modify: e.g. "pytorch/executorch".',
+        required=True,
+    )
+    parser.add_argument(
+        "--pr",
+        type=int,
+        help="Number of the PR in the stack to check and create follow up PR",
+        required=True,
+    )
+    return parser.parse_args()
+
+
+def extract_stack_from_body(pr_body: str) -> List[int]:
+    """Extracts a list of PR numbers from a ghexport-generated PR body.
+
+    The base of the stack is in index 0.
+    """
+
+    # Expected format. The `__->__` could appear on any line. Stop parsing
+    # after the blank line. This would return [1, 2, 3].
+    """
+    Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
+    * #3
+    * __->__ #2
+    * #1
+
+    <PR description details>
+    """
+
+    prs = []
+    for line in pr_body.splitlines():
+        match = re.match(r"\*(?:.*?)? #(\d+)", line)
+        if match:
+            # It's a bullet followed by an integer.
+            prs.append(int(match.group(1)))
+        elif line.strip() == "":
+            # End of the stack.
+            break
+    return list(reversed(prs))
+
+
+def get_pr_stack_from_number(pr_number: int, repo: Repository) -> List[int]:
+    return extract_stack_from_body(repo.get_pull(pr_number).body)
+
+
+def create_prs_for_orig_branch(pr_stack: List[int], repo: Repository):
+    # For the first PR, we want to merge to `main` branch, and we will update
+    # as we go through the stack
+    orig_branch_merge_base = "main"
+    for i in range(len(pr_stack)):
+        pr = repo.get_pull(pr_stack[i])
+        # Check for invariant: For the current PR, it must be gh/user/x/base <- gh/user/x/head
+        assert pr.base.ref.replace("base", "head") == pr.head.ref
+        # The PR we want to create is then "branch_to_merge" <- gh/user/x/orig
+        # gh/user/x/orig is the clean diff between gh/user/x/base <- gh/user/x/head
+        orig_branch_merge_head = pr.base.ref.replace("base", "orig")
+        bot_metadata = f"""\nThis PR was created by the merge bot to help merge the original PR into the base branch.
+ghstack PR number: https://github.com/pytorch/executorch/pull/{pr.number}
+ghstack PR base: https://github.com/pytorch/executorch/tree/{pr.base.ref}
+ghstack PR head: https://github.com/pytorch/executorch/tree/{pr.head.ref}
+Merge bot PR base: https://github.com/pytorch/executorch/tree/{orig_branch_merge_base}
+Merge bot PR head: https://github.com/pytorch/executorch/tree/{orig_branch_merge_head}
+        """
+        repo.create_pull(base=orig_branch_merge_base, head=orig_branch_merge_head, title=pr.title, body=pr.body + bot_metadata)
+        # Advance the base for the next PR
+        orig_branch_merge_base = orig_branch_merge_head
+
+def main():
+    args = parse_args()
+
+    with Github(auth=Auth.Token(os.environ['GITHUB_TOKEN'])) as gh:
+        repo = gh.get_repo(args.repo)
+        create_prs_for_orig_branch(get_pr_stack_from_number(args.pr, repo), repo)
+    
+
+if __name__ == "__main__":
+    main()

.github/workflows/ghstack_land.yml

@@ -0,0 +1,38 @@
+name: Propose to merge ghstack orig PRs to main
+on:
+  pull_request:
+    types: [opened, synchronize, closed]
+    branches:
+      - 'gh/*/[0-9]+/base'
+      - 'kirklandsigntest/*'
+jobs:
+  ghstack_merge_to_main:
+    name: ghstack-land
+    runs-on: ubuntu-22.04
+    permissions:
+      pull-requests: write
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: '0'
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+    - name: Try to merge PR to main
+      run: |
+        pip install ghstack
+        pip install pygithub
+
+        PR_NUMBER=$(echo "$GITHUB_REF" | grep -oE '[0-9]+')
+
+        echo "Checking whether PR $PR_NUMBER is merged"
+        PR_MERGED=$(python -c "from github import Github; g = Github(); repo = g.get_repo('pytorch/executorch'); pr = repo.get_pull(${PR_NUMBER}); print(1 if pr.merged else 0)")
+        if [ "$PR_MERGED" = "0" ]; then
+          echo $"PR $PR_NUMBER is not merged. No op here."
+          exit 0
+        fi
+
+        python .github/scripts/propose_ghstack_orig_pr.py --pr $PR_NUMBER --repo pytorch/executorch
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        GITHUB_REF: ${{ github.ref }}