From 607ed84fcd45f8eef942ebfef33a4b96dd2988e6 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Tue, 16 Dec 2025 09:40:15 -0500 Subject: [PATCH] Update [ghstack-poisoned] --- src/ghstack/checkout.py | 33 +++++++++++++++++++++++ src/ghstack/cli.py | 8 +++++- src/ghstack/test_prelude.py | 13 +++++++++ test/checkout/basic.py.test | 24 +++++++++++++++++ test/checkout/same_base_allows.py.test | 31 ++++++++++++++++++++++ test/checkout/same_base_rejects.py.test | 35 +++++++++++++++++++++++++ 6 files changed, 143 insertions(+), 1 deletion(-) create mode 100644 test/checkout/basic.py.test create mode 100644 test/checkout/same_base_allows.py.test create mode 100644 test/checkout/same_base_rejects.py.test diff --git a/src/ghstack/checkout.py b/src/ghstack/checkout.py index 4767297..b191f71 100644 --- a/src/ghstack/checkout.py +++ b/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) diff --git a/src/ghstack/cli.py b/src/ghstack/cli.py index d27fe65..82c8cbc 100644 --- a/src/ghstack/cli.py +++ b/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, ) diff --git a/src/ghstack/test_prelude.py b/src/ghstack/test_prelude.py index 78461ba..a46d5a9 100644 --- a/src/ghstack/test_prelude.py +++ b/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: diff --git a/test/checkout/basic.py.test b/test/checkout/basic.py.test new file mode 100644 index 0000000..e4e5cb4 --- /dev/null +++ b/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() diff --git a/test/checkout/same_base_allows.py.test b/test/checkout/same_base_allows.py.test new file mode 100644 index 0000000..ead555a --- /dev/null +++ b/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() diff --git a/test/checkout/same_base_rejects.py.test b/test/checkout/same_base_rejects.py.test new file mode 100644 index 0000000..7c44c2d --- /dev/null +++ b/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()