diff --git a/README.md b/README.md index 3e59e5b..74cc6e7 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,18 @@ You can also setup a GitHub action to allow a bot land. [ghstack_land_example](https://github.com/Chillee/ghstack_land_example/blob/main/.github/workflows/ghstack_land.yml) is an end-to-end example of how to do this. +## GitHub native stacked PRs + +When direct mode is enabled with `--direct` or `.github/ghstack_direct`, +ghstack links chains of two or more pull requests using GitHub's native Stack +API. Existing native stacks are reused, and newly submitted pull requests are +appended to the top. Repositories where native stacked PRs are not enabled +continue to use the existing direct-mode behavior. + +GitHub owns a pull request's base branch while it belongs to a native stack. +To reorder a submitted stack or otherwise change an existing PR's base, +unstack it in the GitHub UI before rerunning ghstack. + ## Structure of submitted pull requests Every commit in your local commit stack gets submitted into a separate diff --git a/src/ghstack/github_fake.py b/src/ghstack/github_fake.py index 4871312..1605198 100644 --- a/src/ghstack/github_fake.py +++ b/src/ghstack/github_fake.py @@ -75,6 +75,7 @@ class GitHubState: repositories: Dict[GraphQLId, "Repository"] pull_requests: Dict[GraphQLId, "PullRequest"] + native_stacks: Dict[int, "NativeStack"] # This is very inefficient but whatever issue_comments: Dict[GraphQLId, "IssueComment"] _next_id: int @@ -83,6 +84,7 @@ class GitHubState: _next_issue_comment_full_database_id: Dict[GraphQLId, int] root: "Root" upstream_sh: Optional[ghstack.shell.Shell] + native_stacks_enabled: bool def repository(self, owner: str, name: str) -> "Repository": nameWithOwner = "{}/{}".format(owner, name) @@ -142,11 +144,13 @@ def notify_merged(self, pr_resolved: ghstack.diff.PullRequestResolved) -> None: def __init__(self, upstream_sh: Optional[ghstack.shell.Shell]) -> None: self.repositories = {} self.pull_requests = {} + self.native_stacks = {} self.issue_comments = {} self._next_id = 5000 self._next_pull_request_number = {} self._next_issue_comment_full_database_id = {} self.root = Root() + self.native_stacks_enabled = False # Populate it with the most important repo ;) repo = Repository( @@ -297,6 +301,12 @@ class PullRequestConnection: nodes: List[PullRequest] +@dataclass +class NativeStack: + number: int + pull_requests: List[GitHubNumber] + + class Root: def repository(self, info: GraphQLResolveInfo, owner: str, name: str) -> Repository: return github_state(info).repository(owner, name) @@ -417,11 +427,66 @@ async def _update_pull_async( if "title" in input and input["title"] is not None: pr.title = input["title"] if "base" in input and input["base"] is not None: + if any( + number in stack.pull_requests for stack in state.native_stacks.values() + ): + raise RuntimeError("Native stacks own pull request base branches") pr.baseRefName = input["base"] pr.baseRef = await repo._make_ref_async(state, pr.baseRefName) if "body" in input and input["body"] is not None: pr.body = input["body"] + def _native_stack_payload( + self, state: GitHubState, stack: NativeStack + ) -> Dict[str, Any]: + repo = state.repository("pytorch", "pytorch") + pull_requests = [ + state.pull_request(repo, number) for number in stack.pull_requests + ] + return { + "id": stack.number, + "number": stack.number, + "pull_requests": [{"number": pr.number} for pr in pull_requests], + } + + def _validate_native_stack( + self, state: GitHubState, pull_requests: Sequence[GitHubNumber] + ) -> None: + if len(pull_requests) < 2: + raise RuntimeError("Stack must contain at least two pull requests") + repo = state.repository("pytorch", "pytorch") + prs = [state.pull_request(repo, number) for number in pull_requests] + for lower, upper in zip(prs, prs[1:]): + if upper.baseRefName != lower.headRefName: + raise RuntimeError("Pull requests must form a stack") + + def _create_native_stack( + self, state: GitHubState, pull_requests: Sequence[int] + ) -> Dict[str, Any]: + numbers = [GitHubNumber(number) for number in pull_requests] + self._validate_native_stack(state, numbers) + stacked = { + number + for stack in state.native_stacks.values() + for number in stack.pull_requests + } + if any(number in stacked for number in numbers): + raise RuntimeError("Pull requests are already part of a stack") + repo = state.repository("pytorch", "pytorch") + number = int(state.next_pull_request_number(repo.id)) + stack = NativeStack(number=number, pull_requests=numbers) + state.native_stacks[number] = stack + return self._native_stack_payload(state, stack) + + def _add_to_native_stack( + self, state: GitHubState, stack_number: int, pull_requests: Sequence[int] + ) -> Dict[str, Any]: + stack = state.native_stacks[stack_number] + additions = [GitHubNumber(number) for number in pull_requests] + self._validate_native_stack(state, [stack.pull_requests[-1], *additions]) + stack.pull_requests.extend(additions) + return self._native_stack_payload(state, stack) + async def _create_issue_comment_async( self, owner: str, name: str, comment_id: int, input: CreateIssueCommentInput ) -> CreateIssueCommentPayload: @@ -467,6 +532,21 @@ async def arest(self, method: str, path: str, **kwargs: Any) -> Any: async def _arest_impl(self, method: str, path: str, **kwargs: Any) -> Any: if method == "get": + if m := re.match( + r"^repos/([^/]+)/([^/]+)/stacks(?:\?pull_request=([0-9]+))?$", + path, + ): + if not self.state.native_stacks_enabled: + raise ghstack.github.NotFoundError() + pull_request = ( + GitHubNumber(int(m.group(3))) if m.group(3) is not None else None + ) + stacks = [ + self._native_stack_payload(self.state, stack) + for stack in self.state.native_stacks.values() + if pull_request is None or pull_request in stack.pull_requests + ] + return stacks m = re.match(r"^repos/([^/]+)/([^/]+)/branches/([^/]+)/protection", path) if m: # For now, pretend all branches are not protected @@ -497,6 +577,18 @@ async def _arest_impl(self, method: str, path: str, **kwargs: Any) -> Any: } elif method == "post": + if m := re.match(r"^repos/([^/]+)/([^/]+)/stacks$", path): + if not self.state.native_stacks_enabled: + raise ghstack.github.NotFoundError() + return self._create_native_stack(self.state, kwargs["pull_requests"]) + if m := re.match(r"^repos/([^/]+)/([^/]+)/stacks/([0-9]+)/add$", path): + if not self.state.native_stacks_enabled: + raise ghstack.github.NotFoundError() + return self._add_to_native_stack( + self.state, + int(m.group(3)), + kwargs["pull_requests"], + ) if m := re.match(r"^repos/([^/]+)/([^/]+)/pulls$", path): return await self._create_pull_async( m.group(1), m.group(2), cast(CreatePullRequestInput, kwargs) diff --git a/src/ghstack/submit.py b/src/ghstack/submit.py index a649f03..d52c639 100644 --- a/src/ghstack/submit.py +++ b/src/ghstack/submit.py @@ -1977,38 +1977,6 @@ async def push_updates( all_diffs: Optional[List[DiffMeta]] = None, import_help: bool = True, ) -> None: - # Collect all refspecs into a single batched push. This is being - # tested in production because GitHub may observe base/head ref updates - # out of order when refreshing PR diffs. If that happens, revert this - # block to three grouped pushes in this order: - # 1. base branches - # 2. head/next branches - # 3. orig branches - # Per-refspec force is encoded with the + prefix: - # orig branches: always force-pushed - # head/next branches: force-pushed only with --force flag - # base branches: never force-pushed - # It is VERY important that we preserve base-before-head ordering, - # otherwise GitHub can spuriously think that the user pushed a number - # of patches as part of the PR, when actually they were just from the - # new upstream branch. - all_push_specs: List[str] = [] - - for s in reversed(diffs_to_submit): - for diff, b in s.push_branches: - # Careful! Don't push main. - if b == "orig": - force = True - elif b == "base": - force = False - else: - force = self.force - all_push_specs.append( - push_spec(diff, branch(s.username, s.ghnum, b), force=force) - ) - if all_push_specs: - await self._git_push(all_push_specs) - # Discover orphan PR numbers from the old stack listing. # We search the full local stack for old stack text, then # collect open PRs that aren't being submitted — both above @@ -2049,6 +2017,49 @@ async def push_updates( orphan_above.append(num) break + native_pr_numbers = self._native_stack_pr_numbers( + diffs_to_submit, orphan_above, orphan_below + ) + native_stacks_available = False + native_stack = None + if self.direct and len(native_pr_numbers) >= 2: + native_stacks_available, native_stack = await self._prepare_native_stack( + native_pr_numbers, + diffs_to_submit, + ) + + # Collect all refspecs into a single batched push. This is being + # tested in production because GitHub may observe base/head ref updates + # out of order when refreshing PR diffs. If that happens, revert this + # block to three grouped pushes in this order: + # 1. base branches + # 2. head/next branches + # 3. orig branches + # Per-refspec force is encoded with the + prefix: + # orig branches: always force-pushed + # head/next branches: force-pushed only with --force flag + # base branches: never force-pushed + # It is VERY important that we preserve base-before-head ordering, + # otherwise GitHub can spuriously think that the user pushed a number + # of patches as part of the PR, when actually they were just from the + # new upstream branch. + all_push_specs: List[str] = [] + + for s in reversed(diffs_to_submit): + for diff, b in s.push_branches: + # Careful! Don't push main. + if b == "orig": + force = True + elif b == "base": + force = False + else: + force = self.force + all_push_specs.append( + push_spec(diff, branch(s.username, s.ghnum, b), force=force) + ) + if all_push_specs: + await self._git_push(all_push_specs) + def _update_pr_args(s: DiffMeta) -> Tuple[str, Dict[str, Any], Optional[str]]: assert not s.closed logging.info( @@ -2060,10 +2071,10 @@ def _update_pr_args(s: DiffMeta) -> Tuple[str, Dict[str, Any], Optional[str]]: ) ) base_kwargs = {} - if self.direct: + if self.direct and s.base != s.elab_diff.base_ref: base_kwargs["base"] = s.base else: - assert s.base == s.elab_diff.base_ref + assert self.direct or s.base == s.elab_diff.base_ref stack_desc = self._format_stack( diffs_to_submit, s.number, orphan_above, orphan_below ) @@ -2102,6 +2113,9 @@ async def _update_pr_async(s: DiffMeta) -> None: await _gather_ordered(_update_pr_async(s) for s in reversed(diffs_to_submit)) + if native_stacks_available: + await self._publish_native_stack(native_pr_numbers, native_stack) + # Report what happened def format_url(s: DiffMeta) -> str: return "https://{github_url}/{owner}/{repo}/pull/{number}".format( @@ -2304,6 +2318,123 @@ def assert_eq(a: Any, b: Any) -> None: # ~~~~~~~~~~~~~~~~~~~~~~~~ # Small helpers + def _native_stack_pr_numbers( + self, + diffs_to_submit: Sequence[DiffMeta], + orphan_above: Sequence[GitHubNumber], + orphan_below: Sequence[GitHubNumber], + ) -> List[GitHubNumber]: + top_to_bottom = [ + *orphan_above, + *(s.number for s in diffs_to_submit), + *orphan_below, + ] + seen: Set[GitHubNumber] = set() + result = [] + for number in reversed(top_to_bottom): + if number not in seen: + result.append(number) + seen.add(number) + return result + + def _remote_stack_pr_numbers(self, stack: Any) -> List[GitHubNumber]: + result = [] + for pr in stack.get("pull_requests", []): + number = pr.get("number") if isinstance(pr, dict) else pr + if isinstance(number, int): + result.append(GitHubNumber(number)) + return result + + async def _find_native_stack( + self, pr_numbers: Sequence[GitHubNumber] + ) -> Tuple[bool, Optional[Any]]: + for number in pr_numbers: + try: + stacks = await self.github.aget( + f"repos/{self.repo_owner}/{self.repo_name}/stacks" + f"?pull_request={number}" + ) + except ghstack.github.NotFoundError: + return False, None + except RuntimeError as e: + logging.warning("Failed to query GitHub native stacks: %s", e) + return False, None + if stacks: + return True, stacks[0] + return True, None + + async def _prepare_native_stack( + self, + pr_numbers: List[GitHubNumber], + diffs_to_submit: Sequence[DiffMeta], + ) -> Tuple[bool, Optional[Any]]: + available, stack = await self._find_native_stack(pr_numbers) + if not available or stack is None: + return available, stack + + current = self._remote_stack_pr_numbers(stack) + if current != pr_numbers[: len(current)]: + for prefix_len in range(1, len(current)): + remote_prefix = current[:prefix_len] + remote_suffix = current[prefix_len:] + if remote_suffix == pr_numbers[: len(remote_suffix)] and not set( + remote_prefix + ).intersection(pr_numbers): + pr_numbers[:0] = remote_prefix + break + else: + raise RuntimeError( + "The GitHub native stack has a different pull request order. " + "Native stacks own their pull requests' base branches, so ghstack " + "cannot reorder it directly. Unstack it in the GitHub UI, then " + "rerun ghstack." + ) + + current_numbers = set(current) + for diff_meta in diffs_to_submit: + if ( + diff_meta.number in current_numbers + and diff_meta.base != diff_meta.elab_diff.base_ref + ): + raise RuntimeError( + f"GitHub native stack #{stack['number']} owns the base branch " + f"for PR #{diff_meta.number}. Unstack it in the GitHub UI, " + "then rerun ghstack." + ) + return available, stack + + async def _publish_native_stack( + self, pr_numbers: List[GitHubNumber], stack: Optional[Any] + ) -> None: + path = f"repos/{self.repo_owner}/{self.repo_name}/stacks" + try: + if stack is None: + result = await self.github.apost( + path, pull_requests=[int(n) for n in pr_numbers] + ) + logging.info( + "Created GitHub native stack #%s", result.get("number", "unknown") + ) + return + + current = self._remote_stack_pr_numbers(stack) + delta = pr_numbers[len(current) :] + if not delta: + return + result = await self.github.apost( + f"{path}/{stack['number']}/add", + pull_requests=[int(n) for n in delta], + ) + logging.info( + "Added %d pull request(s) to GitHub native stack #%s", + len(delta), + result.get("number", stack["number"]), + ) + except ghstack.github.NotFoundError: + logging.info("GitHub native stacks are not enabled for this repository") + except RuntimeError as e: + logging.warning("Failed to update GitHub native stack: %s", e) + # TODO: do the tree formatting minigame # Main things: # - need to express some tree structure diff --git a/src/ghstack/test_prelude.py b/src/ghstack/test_prelude.py index 69ea201..5dc1929 100644 --- a/src/ghstack/test_prelude.py +++ b/src/ghstack/test_prelude.py @@ -69,6 +69,8 @@ "get_sh", "get_upstream_sh", "get_github", + "enable_native_stacks", + "get_native_stacks", "get_pr_reviewers", "get_pr_labels", "tick", @@ -461,6 +463,15 @@ def get_github() -> "ghstack.github_fake.FakeGitHubEndpoint": return github +def enable_native_stacks() -> None: + get_github().state.native_stacks_enabled = True + + +def get_native_stacks() -> List[List[int]]: + stacks = get_github().state.native_stacks.values() + return [[int(number) for number in stack.pull_requests] for stack in stacks] + + def get_pr_reviewers(pr_number: int) -> List[str]: """Get the reviewers for a PR number.""" github = get_github() diff --git a/test/submit/native_stack.py.test b/test/submit/native_stack.py.test new file mode 100644 index 0000000..2abeac7 --- /dev/null +++ b/test/submit/native_stack.py.test @@ -0,0 +1,25 @@ +from ghstack.test_prelude import * + +await init_test() +enable_native_stacks() + +await commit("A") +await commit("B") +A, B = await gh_submit("Initial") + +if is_direct(): + assert_eq(get_native_stacks(), [[500, 501]]) + + await gh_submit("No changes") + assert_eq(get_native_stacks(), [[500, 501]]) + + await commit("C") + A2, B2, C = await gh_submit("Append") + assert_eq(A.number, A2.number) + assert_eq(B.number, B2.number) + assert_eq(C.number, 503) + assert_eq(get_native_stacks(), [[500, 501, 503]]) +else: + assert_eq(get_native_stacks(), []) + +ok() diff --git a/test/submit/native_stack_partial_merge.py.test b/test/submit/native_stack_partial_merge.py.test new file mode 100644 index 0000000..389c90e --- /dev/null +++ b/test/submit/native_stack_partial_merge.py.test @@ -0,0 +1,31 @@ +from ghstack.test_prelude import * +from ghstack.github_fake import GitHubNumber + +await init_test() +enable_native_stacks() + +await commit("A") +await commit("B") +A, B = await gh_submit("Initial") + +if is_direct(): + github = get_github() + repo = github.state.repository("pytorch", "pytorch") + pr_a = github.state.pull_request(repo, GitHubNumber(A.number)) + pr_b = github.state.pull_request(repo, GitHubNumber(B.number)) + pr_a.closed = True + pr_b.baseRefName = "main" + pr_b.baseRef = await repo._make_ref_async(github.state, "main") + + await checkout(B) + await git("rebase", "--onto", "origin/main", A.orig) + await commit("C") + B2, C = await gh_submit("Append after partial merge") + + assert_eq(B.number, B2.number) + assert_eq(C.number, 503) + assert_eq(get_native_stacks(), [[500, 501, 503]]) +else: + assert_eq(get_native_stacks(), []) + +ok() diff --git a/test/submit/native_stack_reorder.py.test b/test/submit/native_stack_reorder.py.test new file mode 100644 index 0000000..d013972 --- /dev/null +++ b/test/submit/native_stack_reorder.py.test @@ -0,0 +1,25 @@ +from ghstack.test_prelude import * + +await init_test() +enable_native_stacks() + +await commit("A") +await commit("B") +A, B = await gh_submit("Initial") + +if is_direct(): + await checkout(GitCommitHash("HEAD~~")) + await cherry_pick(B) + await cherry_pick(A) + await assert_expected_raises_inline( + RuntimeError, + gh_submit, + """\ +The GitHub native stack has a different pull request order. Native stacks own their pull requests' base branches, so ghstack cannot reorder it directly. Unstack it in the GitHub UI, then rerun ghstack.""", + "Reorder", + ) + assert_eq(get_native_stacks(), [[500, 501]]) +else: + assert_eq(get_native_stacks(), []) + +ok()