-
Notifications
You must be signed in to change notification settings - Fork 943
Expand file tree
/
Copy pathpropose_ghstack_orig_pr.py
More file actions
155 lines (129 loc) · 5.29 KB
/
propose_ghstack_orig_pr.py
File metadata and controls
155 lines (129 loc) · 5.29 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
# 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(
"--ref",
type=str,
help="Ref fo PR in the stack to check and create corresponding 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 = []
ghstack_begin = (
"Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):"
)
ghstack_begin_seen = False
for line in pr_body.splitlines():
if ghstack_begin in line:
ghstack_begin_seen = True
if not ghstack_begin_seen:
continue
match = re.match(r"\*(?:.*?)? #(\d+)", line)
if match:
# It's a bullet followed by an integer.
prs.append(int(match.group(1)))
return list(reversed(prs))
def get_pr_stack_from_number(ref: str, repo: Repository) -> List[int]:
if ref.isnumeric():
pr_number = int(ref)
else:
branch_name = ref.replace("refs/heads/", "")
pr_number = repo.get_branch(branch_name).commit.get_pulls()[0].number
pr_stack = extract_stack_from_body(repo.get_pull(pr_number).body)
if not pr_stack:
raise Exception(
f"Could not find PR stack in body of ref. "
+ "Please make sure that the PR was created with ghstack."
)
return pr_stack
def get_differential_revision(pr, repo: Repository) -> str:
body = repo.get_pull(pr.number).body
matches = re.findall(r"Differential Revision: .*", body)
count = len(matches)
if count == 1:
# If there's more than one Differential Revision, let's just return empty
# so that we can disambiguate manually.
return matches[0]
return ""
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])
if not pr.is_merged():
print("The PR (and stack above) is not merged yet, skipping")
return
# 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")
differential_revision_text = get_differential_revision(pr, repo)
bot_metadata = f"""This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: https://github.com/pytorch/executorch/pull/{pr.number} by @{pr.user.login}
^ Please use this as the source of truth for the PR details, comments, and reviews
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}
{differential_revision_text}
@diff-train-skip-merge"""
existing_orig_pr = repo.get_pulls(
head="pytorch:" + orig_branch_merge_head,
base=orig_branch_merge_base,
state="all",
)
if existing_orig_pr.totalCount > 0 and existing_orig_pr[0].title == pr.title:
print(
f"PR for {orig_branch_merge_head} already exists {existing_orig_pr[0]}"
)
# We don't need to create/edit because the head PR is merged and orig is finalized.
else:
repo.create_pull(
base=orig_branch_merge_base,
head=orig_branch_merge_head,
title=pr.title,
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.ref, repo), repo)
if __name__ == "__main__":
main()