Update

[ghstack-poisoned]

Author: ezyang

src/ghstack/checkout.py

@@ -13,6 +13,7 @@ def main(
     github: ghstack.github.GitHubEndpoint,
     sh: ghstack.shell.Shell,
     remote_name: str,
+    same_base: bool = False,
 ) -> None:
 
     params = ghstack.github_utils.parse_pull_request(
@@ -27,5 +28,37 @@ def main(
 
     # TODO: Handle remotes correctly too (so this subsumes hub)
 
+    # If --same-base is specified, check if checkout would change the merge-base
+    if same_base:
+        # Get the default branch name from the repo
+        repo_info = ghstack.github_utils.get_github_repo_info(
+            github=github,
+            sh=sh,
+            repo_owner=params["owner"],
+            repo_name=params["name"],
+            github_url=params["github_url"],
+            remote_name=remote_name,
+        )
+        default_branch = repo_info["default_branch"]
+        default_branch_ref = f"{remote_name}/{default_branch}"
+
+        # Get current merge-base with default branch
+        current_base = sh.git("merge-base", default_branch_ref, "HEAD")
+    else:
+        current_base = None
+        default_branch_ref = None
+
     sh.git("fetch", "--prune", remote_name)
+
+    # If --same-base is specified, check what the new merge-base would be
+    if same_base:
+        target_ref = remote_name + "/" + orig_ref
+        new_base = sh.git("merge-base", default_branch_ref, target_ref)
+
+        if current_base != new_base:
+            raise RuntimeError(
+                f"Checkout would change merge-base from {current_base[:8]} to {new_base[:8]}, "
+                f"aborting due to --same-base flag"
+            )
+
     sh.git("checkout", remote_name + "/" + orig_ref)

src/ghstack/cli.py

@@ -120,8 +120,13 @@ def action(close: bool, pull_request: str) -> None:
 
 
 @main.command("checkout")
+@click.option(
+    "--same-base",
+    is_flag=True,
+    help="Only checkout if merge-base with main branch would remain the same",
+)
 @click.argument("pull_request", metavar="PR")
-def checkout(pull_request: str) -> None:
+def checkout(same_base: bool, pull_request: str) -> None:
     """
     Checkout a PR
     """
@@ -131,6 +136,7 @@ def checkout(pull_request: str) -> None:
             github=github,
             sh=shell,
             remote_name=config.remote_name,
+            same_base=same_base,
         )
 
 

src/ghstack/test_prelude.py

@@ -12,6 +12,7 @@
 
 from expecttest import assert_expected_inline
 
+import ghstack.checkout
 import ghstack.cherry_pick
 
 import ghstack.github
@@ -32,6 +33,7 @@
     "gh_land",
     "gh_unlink",
     "gh_cherry_pick",
+    "gh_checkout",
     "GitCommitHash",
     "checkout",
     "amend",
@@ -251,6 +253,17 @@ def gh_cherry_pick(pull_request: str, stack: bool = False) -> None:
     )
 
 
+def gh_checkout(pull_request: str, same_base: bool = False) -> None:
+    self = CTX
+    return ghstack.checkout.main(
+        pull_request=pull_request,
+        github=self.github,
+        sh=self.sh,
+        remote_name="origin",
+        same_base=same_base,
+    )
+
+
 def write_file_and_add(filename: str, contents: str) -> None:
     self = CTX
     with self.sh.open(filename, "w") as f:

test/checkout/basic.py.test

@@ -0,0 +1,24 @@
+from ghstack.test_prelude import *
+
+init_test()
+
+# Create a PR to checkout
+commit("A")
+(A,) = gh_submit("Initial commit")
+
+# Move to master and create another commit
+git("checkout", "master")
+commit("B")
+
+# Verify we're on master with commit B
+current_log = git("log", "--oneline", "-n", "1")
+assert "Commit B" in current_log
+
+# Checkout the PR
+gh_checkout(f"https://github.com/pytorch/pytorch/pull/{A.number}")
+
+# After checkout, we should be on the PR commit
+current_log = git("log", "--oneline", "-n", "1")
+assert "Commit A" in current_log
+
+ok()

test/checkout/same_base_allows.py.test

@@ -0,0 +1,31 @@
+from ghstack.test_prelude import *
+
+init_test()
+
+# Create two PRs in a stack - they'll have the same base
+commit("A")
+commit("B")
+diffs = gh_submit("Stack of two commits")
+
+# Should have two PRs
+assert len(diffs) == 2
+A = diffs[0]  # First commit (A)
+B = diffs[1]  # Second commit (B)
+
+# Both PRs should have the same merge-base with master (initial commit)
+# Checkout PR A
+gh_checkout(f"https://github.com/pytorch/pytorch/pull/{A.number}")
+
+# Verify we're on PR A
+current_log = git("log", "--oneline", "-n", "1")
+assert "Commit A" in current_log
+
+# Now checkout PR B with --same-base
+# Since both have the same merge-base (initial commit), this should succeed
+gh_checkout(f"https://github.com/pytorch/pytorch/pull/{B.number}", same_base=True)
+
+# Verify we successfully checked out PR B
+current_log = git("log", "--oneline", "-n", "1")
+assert "Commit B" in current_log
+
+ok()

test/checkout/same_base_rejects.py.test

@@ -0,0 +1,35 @@
+import pytest
+from ghstack.test_prelude import *
+
+init_test()
+
+# Create first PR based on initial master
+commit("A")
+(A,) = gh_submit("First PR")
+
+# Go back to master and advance it
+git("checkout", "master")
+commit("B")
+git("push", "origin", "master")
+
+# Create second PR based on new master (different merge-base)
+commit("C")
+(C,) = gh_submit("Second PR")
+
+# Checkout first PR
+gh_checkout(f"https://github.com/pytorch/pytorch/pull/{A.number}")
+
+# Verify we're on PR A
+current_log = git("log", "--oneline", "-n", "1")
+assert "Commit A" in current_log
+
+# Try to checkout second PR with --same-base
+# This should fail because merge-base would change from initial commit to commit B
+with pytest.raises(RuntimeError, match="would change merge-base"):
+    gh_checkout(f"https://github.com/pytorch/pytorch/pull/{C.number}", same_base=True)
+
+# Verify we're still on PR A (checkout was aborted)
+current_log = git("log", "--oneline", "-n", "1")
+assert "Commit A" in current_log
+
+ok()