Add native GitHub Stacks support (--github-stacks)

README.md

@@ -41,6 +41,25 @@ patch update itself is included directly in the prompt.  For existing PRs the
 patch update is an interdiff against the previously submitted PR patch, not the
 whole PR patch.  Explicit `ghstack -m MESSAGE` still overrides this.
 
+## Native GitHub Stacks
+
+If your repository has [GitHub Stacks](https://github.github.com/gh-stack/)
+enabled (currently a private preview), ghstack can link the PRs it submits
+into a native stack, which gets you the stack navigation UI, stack-aware
+branch protection, and partial-stack merges.  Enable it with:
+
+```
+ghstack config github_stacks true
+```
+
+or per-invocation with `ghstack --direct --github-stacks`.  This requires
+`--direct` mode, because GitHub requires each PR in a stack to target the
+head branch of the PR below it, which the classic base-branch layout does
+not satisfy.  ghstack keeps the native stack in sync on resubmit: new
+commits on top are appended, and restructures (inserts, reorders) dissolve
+and recreate the stack, since the API has no reorder operation.  Stack
+linking failures are never fatal; the PRs themselves are always submitted.
+
 ## How to use
 
 Make sure you have write permission to the repo you're opening PR with.

src/ghstack/cli.py

@@ -110,6 +110,13 @@ def cli_context(
     default=None,
     help="Create stack that directly merges into main",
 )
+@click.option(
+    "--github-stacks/--no-github-stacks",
+    "github_stacks_opt",
+    is_flag=True,
+    default=None,
+    help="Link submitted PRs into a native GitHub stack (requires --direct)",
+)
 @click.option(
     "--base",
     "-B",
@@ -141,6 +148,7 @@ def main(
     short: bool,
     force: bool,
     direct_opt: Optional[bool],
+    github_stacks_opt: Optional[bool],
     no_skip: bool,
     draft: bool,
     base: Optional[str],
@@ -166,6 +174,7 @@ def main(
             base=base,
             stack=stack,
             direct_opt=direct_opt,
+            github_stacks_opt=github_stacks_opt,
             reviewer=reviewer,
             label=label,
         )
@@ -505,6 +514,14 @@ def status(pull_request: str) -> None:
     is_flag=True,
     help="Create stack that directly merges into main",
 )
+@click.option(
+    "--github-stacks/--no-github-stacks",
+    "github_stacks_opt",
+    default=None,
+    is_flag=True,
+    help="Link submitted PRs into a native GitHub stack (requires --direct and "
+    "a repository with GitHub Stacks enabled; overrides .ghstackrc setting)",
+)
 @click.option(
     "--no-fetch",
     is_flag=True,
@@ -523,6 +540,7 @@ def submit(
     no_skip: bool,
     draft: bool,
     direct_opt: Optional[bool],
+    github_stacks_opt: Optional[bool],
     base: Optional[str],
     revs: Tuple[str, ...],
     stack: bool,
@@ -553,6 +571,11 @@ def submit(
                     revs=revs,
                     stack=stack,
                     direct_opt=direct_opt,
+                    github_stacks=(
+                        github_stacks_opt
+                        if github_stacks_opt is not None
+                        else config.github_stacks
+                    ),
                     reviewer=reviewer if reviewer is not None else config.reviewer,
                     label=label if label is not None else config.label,
                     no_fetch=no_fetch,

src/ghstack/config.py

@@ -113,6 +113,9 @@ def get_gh_cli_credentials(
         ("label", Optional[str]),
         # Command to generate a per-PR update description from diff contents
         ("automsg", Optional[str]),
+        # Link submitted PRs into a native GitHub stack (requires direct mode
+        # and a repository with GitHub Stacks enabled)
+        ("github_stacks", bool),
     ],
 )
 
@@ -334,6 +337,11 @@ def read_config(
     else:
         automsg = None
 
+    if config.has_option("ghstack", "github_stacks"):
+        github_stacks = config.getboolean("ghstack", "github_stacks")
+    else:
+        github_stacks = False
+
     if write_back:
         with open(config_path, "w") as f:
             config.write(f)
@@ -352,6 +360,7 @@ def read_config(
         reviewer=reviewer,
         label=label,
         automsg=automsg,
+        github_stacks=github_stacks,
     )
     logging.debug(f"conf = {conf}")
     return conf

src/ghstack/github_fake.py

@@ -71,10 +71,20 @@
 )
 
 
+# A native GitHub stack (the REST resource behind `gh stack`).  Stack
+# numbers are allocated from the same sequence as PR numbers.
+@dataclass
+class PullRequestStack:
+    number: GitHubNumber
+    _repository: GraphQLId
+    pull_request_numbers: List[GitHubNumber]
+
+
 # The "database" for our mock instance
 class GitHubState:
     repositories: Dict[GraphQLId, "Repository"]
     pull_requests: Dict[GraphQLId, "PullRequest"]
+    stacks: Dict[GitHubNumber, "PullRequestStack"]
     # This is very inefficient but whatever
     issue_comments: Dict[GraphQLId, "IssueComment"]
     _next_id: int
@@ -142,6 +152,7 @@ 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.stacks = {}
         self.issue_comments = {}
         self._next_id = 5000
         self._next_pull_request_number = {}
@@ -417,6 +428,14 @@ 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:
+            # Real GitHub rejects any PATCH with a base field on a PR that
+            # is part of a native stack, even if the value is unchanged
+            for stack in state.stacks.values():
+                if number in stack.pull_request_numbers:
+                    raise RuntimeError(
+                        "Cannot change the base branch because the pull "
+                        "request is part of a stack (HTTP 422)"
+                    )
             pr.baseRefName = input["base"]
             pr.baseRef = await repo._make_ref_async(state, pr.baseRefName)
         if "body" in input and input["body"] is not None:
@@ -451,6 +470,79 @@ async def _update_issue_comment_async(
         if (r := input.get("body")) is not None:
             comment.body = r
 
+    def _dump_stack(self, owner: str, name: str, stack: PullRequestStack) -> Any:
+        state = self.state
+        repo = state.repository(owner, name)
+        prs = [state.pull_request(repo, n) for n in stack.pull_request_numbers]
+        return {
+            "number": stack.number,
+            "base": {"ref": prs[0].baseRefName},
+            "open": True,
+            "pull_requests": [
+                {
+                    "number": pr.number,
+                    "state": "closed" if pr.closed else "open",
+                    "head": {"ref": pr.headRefName},
+                    "base": {"ref": pr.baseRefName},
+                }
+                for pr in prs
+            ],
+        }
+
+    # Mimics the base-chaining rule the real server enforces on both
+    # create and add: each PR's base ref must be the previous PR's head ref
+    def _validate_stack_chain(
+        self, repo: Repository, numbers: List[GitHubNumber], prev_head: Optional[str]
+    ) -> None:
+        state = self.state
+        for number in numbers:
+            for stack in state.stacks.values():
+                if number in stack.pull_request_numbers:
+                    raise RuntimeError(
+                        f"Pull request #{number} is already in a stack (HTTP 422)"
+                    )
+            pr = state.pull_request(repo, number)
+            if prev_head is not None and pr.baseRefName != prev_head:
+                raise RuntimeError(
+                    "Pull requests must form a stack, where each PR's base ref "
+                    "is the previous PR's head ref (HTTP 422)"
+                )
+            prev_head = pr.headRefName
+
+    async def _create_stack_async(
+        self, owner: str, name: str, pull_requests: List[int]
+    ) -> Any:
+        state = self.state
+        repo = state.repository(owner, name)
+        numbers = [GitHubNumber(n) for n in pull_requests]
+        if len(numbers) < 2:
+            raise RuntimeError("A stack must contain at least two PRs (HTTP 422)")
+        self._validate_stack_chain(repo, numbers, None)
+        number = state.next_pull_request_number(repo.id)
+        stack = PullRequestStack(
+            number=number,
+            _repository=repo.id,
+            pull_request_numbers=numbers,
+        )
+        state.stacks[number] = stack
+        return self._dump_stack(owner, name, stack)
+
+    async def _add_to_stack_async(
+        self,
+        owner: str,
+        name: str,
+        stack_number: GitHubNumber,
+        pull_requests: List[int],
+    ) -> Any:
+        state = self.state
+        repo = state.repository(owner, name)
+        stack = state.stacks[stack_number]
+        numbers = [GitHubNumber(n) for n in pull_requests]
+        top = state.pull_request(repo, stack.pull_request_numbers[-1])
+        self._validate_stack_chain(repo, numbers, top.headRefName)
+        stack.pull_request_numbers.extend(numbers)
+        return self._dump_stack(owner, name, stack)
+
     # NB: This may have a payload, but we don't
     # use it so I didn't bother constructing it.
     async def _set_default_branch_async(
@@ -495,6 +587,18 @@ async def _arest_impl(self, method: str, path: str, **kwargs: Any) -> Any:
                     "id": comment.fullDatabaseId,
                     "body": comment.body,
                 }
+            if m := re.match(
+                r"^repos/([^/]+)/([^/]+)/stacks\?pull_request=(\d+)$", path
+            ):
+                pr_number = int(m.group(3))
+                return [
+                    self._dump_stack(m.group(1), m.group(2), stack)
+                    for stack in self.state.stacks.values()
+                    if pr_number in stack.pull_request_numbers
+                ]
+            if m := re.match(r"^repos/([^/]+)/([^/]+)/stacks/(\d+)$", path):
+                stack = self.state.stacks[GitHubNumber(int(m.group(3)))]
+                return self._dump_stack(m.group(1), m.group(2), stack)
 
         elif method == "post":
             if m := re.match(r"^repos/([^/]+)/([^/]+)/pulls$", path):
@@ -518,6 +622,20 @@ async def _arest_impl(self, method: str, path: str, **kwargs: Any) -> Any:
                 reviewers = kwargs.get("reviewers", [])
                 pr.reviewers.extend(reviewers)
                 return {}
+            if m := re.match(r"^repos/([^/]+)/([^/]+)/stacks$", path):
+                return await self._create_stack_async(
+                    m.group(1), m.group(2), kwargs["pull_requests"]
+                )
+            if m := re.match(r"^repos/([^/]+)/([^/]+)/stacks/(\d+)/add$", path):
+                return await self._add_to_stack_async(
+                    m.group(1),
+                    m.group(2),
+                    GitHubNumber(int(m.group(3))),
+                    kwargs["pull_requests"],
+                )
+            if m := re.match(r"^repos/([^/]+)/([^/]+)/stacks/(\d+)/unstack$", path):
+                del self.state.stacks[GitHubNumber(int(m.group(3)))]
+                return None
             if m := re.match(r"^repos/([^/]+)/([^/]+)/issues/([^/]+)/labels", path):
                 # Handle adding labels
                 state = self.state

src/ghstack/github_real.py

@@ -224,6 +224,8 @@ async def arest(self, method: str, path: str, **kwargs: Any) -> Any:
 
             async with getattr(session, method)(url, **request_kwargs) as resp:
                 logging.debug("%s response status: %s", log_prefix, resp.status)
+                if resp.status == 204:
+                    return None
                 try:
                     r = await resp.json()
                 except (aiohttp.ContentTypeError, ValueError):