jj-stack 0.1.0__py3-none-any.whl

jj_stack/commands/checkout.py

@@ -0,0 +1,668 @@
+"""Check out an existing stack of pull requests.
+
+Use `--pull-request` to select a GitHub pull request, `--revset` to select a locally tracked head,
+or `--pick` to choose from local and GitHub stacks in an interactive numbered list. When the
+selected stack's submitted commits are not present locally, the command fetches them
+automatically. It records which pull request belongs to each local change, then runs `jj edit` on
+the selected change.
+
+`jj-stack` changes the working copy only after validating the entire stack and saving any new
+pull request links. `checkout` does not rebase changes or modify GitHub. To create a new change
+on top of the checked-out change, run `jj new` afterward.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from collections import Counter
+from dataclasses import dataclass
+from pathlib import Path
+
+import jj_stack.console as console
+import jj_stack.ui as ui
+from jj_stack.bootstrap import CommandContext, bootstrap_context
+from jj_stack.errors import CliError, UsageError
+from jj_stack.github.client import GithubClient, GithubClientError, build_github_client
+from jj_stack.github.pr_refs import parse_repo_pr_reference
+from jj_stack.github.resolution import (
+    GithubRepoAddress,
+    require_github_repo,
+    select_submit_remote,
+)
+from jj_stack.github.stack_availability import github_stacks_unavailable_error
+from jj_stack.jj.cli_args import JjCliArgs
+from jj_stack.jj.client import JjClient, UnsupportedStackError
+from jj_stack.models.github import GithubPR, GithubStack
+from jj_stack.models.stack import LocalCommit, LocalStack
+from jj_stack.models.tracking import PRIdentity, SubmittedBaseline, TrackingState
+from jj_stack.pr_branch_namespace import current_pr_branch_namespace, pr_branch_matches_change
+from jj_stack.stack.pr_branches import prepare_visible_pr_snapshots
+from jj_stack.stack.pr_facts import duplicate_pr_claim_change_ids
+from jj_stack.stack.repo import observe_repo_paths
+from jj_stack.stack.selected import select_stack_path
+from jj_stack.stack.status import status_preparation_cli_error
+from jj_stack.state.operation_lock import acquire_operation_lock
+
+HELP = "Check out an existing stack of pull requests"
+
+
+@dataclass(frozen=True, slots=True)
+class CheckoutResult:
+    """Outcome of attaching one exact pull-request stack."""
+
+    adopted_count: int
+    fetched_tip_commit: str | None
+    stack: LocalStack
+
+
+@dataclass(frozen=True, slots=True)
+class CheckoutPickerChoice:
+    """One local or GitHub stack offered by the interactive picker."""
+
+    details: tuple[str, ...]
+    heading: str
+    pr: str | None = None
+    revset: str | None = None
+
+
+def checkout(
+    *,
+    cli_args: JjCliArgs,
+    debug: bool,
+    pick: bool,
+    pr: str | None,
+    repo: Path | None,
+    revset: str | None,
+) -> int:
+    """CLI entrypoint for `checkout`."""
+
+    context = bootstrap_context(repo=repo, cli_args=cli_args, debug=debug)
+    if pick:
+        choice = asyncio.run(_pick_stack(context))
+        pr = choice.pr
+        revset = choice.revset
+    with acquire_operation_lock(context.state_store.require_writable(), command="checkout"):
+        result = asyncio.run(
+            _run_checkout_async(
+                context=context,
+                pr_reference=pr,
+                revset=revset,
+            )
+        )
+        if result.stack.changes:
+            if result.adopted_count:
+                prepare_visible_pr_snapshots(
+                    jj_client=context.jj_client,
+                    state=context.state_store.load(),
+                )
+            context.jj_client.edit_commit(result.stack.head.commit_id)
+    if result.fetched_tip_commit is not None:
+        console.output(ui.prefixed_line("Fetched tip commit: ", result.fetched_tip_commit))
+    if result.adopted_count:
+        noun = "PR" if result.adopted_count == 1 else "PRs"
+        console.output(f"Updated local tracking for {result.adopted_count} {noun}.")
+    elif result.stack.changes:
+        console.output("Local tracking is already up to date for this stack.")
+    else:
+        console.output("The selected stack has no changes to check out.")
+    if result.stack.changes:
+        console.output(
+            t"Working copy now edits {ui.change_id(result.stack.head.change_id)} "
+            t"({result.stack.head.subject})."
+        )
+    return 0
+
+
+async def _run_checkout_async(
+    *,
+    context: CommandContext,
+    pr_reference: str | None,
+    revset: str | None,
+) -> CheckoutResult:
+    if pr_reference is not None and revset is not None:
+        raise UsageError(
+            t"{ui.cmd('checkout')} accepts at most one selector: "
+            t"{ui.cmd('--pull-request')} or {ui.cmd('--revset')}."
+        )
+    if pr_reference is None:
+        return _checkout_saved_stack(context=context, revset=revset)
+    return await _checkout_pr_stack(
+        context=context,
+        pr_reference=pr_reference,
+    )
+
+
+def _checkout_saved_stack(
+    *,
+    context: CommandContext,
+    revset: str | None,
+) -> CheckoutResult:
+    client = context.jj_client
+    state = context.state_store.load()
+    stack = select_stack_path(
+        jj_client=client,
+        revset=revset,
+        state=state,
+    ).stack
+    incomplete = tuple(
+        change for change in stack.changes if state.pr_identities.get(change.change_id) is None
+    )
+    if incomplete:
+        raise CliError(
+            t"jj-stack has no saved pull request for some changes in this stack: "
+            t"{ui.join(ui.change_id, (change.change_id for change in incomplete))}.",
+            hint=t"Attach it with {ui.cmd('checkout --pull-request PR')}.",
+        )
+    return CheckoutResult(adopted_count=0, fetched_tip_commit=None, stack=stack)
+
+
+async def _checkout_pr_stack(
+    *,
+    context: CommandContext,
+    pr_reference: str,
+) -> CheckoutResult:
+    client = context.jj_client
+    state = context.state_store.load()
+    remote = select_submit_remote(client.list_git_remotes())
+    repo = require_github_repo(remote)
+    pr_number = parse_repo_pr_reference(
+        reference=pr_reference,
+        github_repo=repo,
+    )
+    async with build_github_client(repo=repo) as github_client:
+        top_pr = await _load_pr(
+            github_client=github_client,
+            pr_number=pr_number,
+        )
+        _validate_same_repo_managed_pr(
+            pr=top_pr,
+            repo=repo,
+        )
+        top_head_sha = _require_pr_head_sha(top_pr)
+        observed_top = client.list_remote_branches(
+            remote=remote.name,
+            patterns=(f"refs/heads/{top_pr.head.ref}",),
+        ).get(top_pr.head.ref)
+        if observed_top != top_head_sha:
+            raise CliError(
+                t"PR #{pr_number} and remote branch "
+                t"{ui.bookmark(top_pr.head.ref)} no longer identify the same commit."
+            )
+        prs = await _load_pr_chain(
+            github_client=github_client,
+            repo=repo,
+            top=top_pr,
+        )
+
+        for pr in reversed(prs):
+            _reject_locally_rewritten_change(
+                client=client,
+                head_sha=_require_pr_head_sha(pr),
+                pr_number=pr.number,
+                remote_name=remote.name,
+            )
+        matches = client.query_commits_by_ids((top_head_sha,))
+        fetched = not matches
+        if fetched:
+            client.fetch_remote(
+                remote=remote.name,
+            )
+            with client.import_remote_pr_branch_ref(
+                remote=remote.name,
+                branch=top_pr.head.ref,
+                expected_target=top_head_sha,
+            ) as imported:
+                _require_branch_matches_change(
+                    branch=top_pr.head.ref,
+                    change=imported,
+                )
+                stack = _discover_checkout_stack(
+                    client=client,
+                    commit_id=imported.commit_id,
+                    state=state,
+                )
+        else:
+            _require_branch_matches_change(
+                branch=top_pr.head.ref,
+                change=matches[0],
+            )
+            stack = _discover_checkout_stack(
+                client=client,
+                commit_id=top_head_sha,
+                state=state,
+            )
+
+        adopted_count = _save_checkout_tracking(
+            context=context,
+            prs=prs,
+            remote_name=remote.name,
+            repo=repo,
+            stack=stack,
+            state=state,
+        )
+    return CheckoutResult(
+        adopted_count=adopted_count,
+        fetched_tip_commit=(top_head_sha if fetched else None),
+        stack=stack,
+    )
+
+
+def _reject_locally_rewritten_change(
+    *,
+    client: JjClient,
+    head_sha: str,
+    pr_number: int,
+    remote_name: str,
+) -> None:
+    """Reject a submitted snapshot that disagrees with a visible local change.
+
+    The remote commit's change ID is read without creating a ref. On a fresh checkout that costs
+    one extra object fetch; reading the same header inside the import primitive would instead
+    give that shared primitive a second policy path. The check also covers an already visible
+    submitted snapshot, where editing it would silently choose against the rewritten local change.
+    """
+
+    change_id = client.read_remote_git_change_id(
+        remote=remote_name,
+        commit_id=head_sha,
+    )
+    if change_id is None:
+        return
+    # `change_id()` rather than an exact symbol: it tolerates a change that is already
+    # divergent, which is the state this check exists to keep the tool out of.
+    local_commits = client.query_commits(f"change_id({change_id})")
+    if not any(commit.commit_id != head_sha for commit in local_commits):
+        return
+    if len(local_commits) > 1:
+        raise CliError(
+            t"Change {ui.change_id(change_id)} already has more than one visible commit "
+            t"here, so PR #{pr_number} cannot be attached to one of them.",
+            hint=t"Inspect them with {ui.cmd('jj log -r')} "
+            t"{ui.revset(f'change_id({change_id})')}, abandon the copies you do not want, "
+            t"then retry.",
+        )
+    raise CliError(
+        t"Change {ui.change_id(change_id)} is already here at a different commit than "
+        t"PR #{pr_number}'s head, so checkout cannot choose between them.",
+        hint=t"Attach the pull request to the local change with "
+        t"{ui.cmd(f'jj-stack relink {pr_number} {change_id}')}.",
+    )
+
+
+def _discover_checkout_stack(
+    *,
+    client: JjClient,
+    commit_id: str,
+    state: TrackingState,
+) -> LocalStack:
+    """Resolve the PR stack, translating shape failures into repair guidance."""
+
+    try:
+        return select_stack_path(
+            jj_client=client,
+            revset=commit_id,
+            state=state,
+        ).stack
+    except UnsupportedStackError as error:
+        raise status_preparation_cli_error(error) from error
+
+
+async def _load_pr(
+    *,
+    github_client: GithubClient,
+    pr_number: int,
+) -> GithubPR:
+    try:
+        return await github_client.get_pr(pr_number=pr_number)
+    except GithubClientError as error:
+        raise CliError(f"Could not load pull request #{pr_number}") from error
+
+
+async def _load_pr_chain(
+    *,
+    github_client: GithubClient,
+    repo: GithubRepoAddress,
+    top: GithubPR,
+) -> tuple[GithubPR, ...]:
+    """Walk exact managed base refs from the selected PR to trunk."""
+
+    namespace = current_pr_branch_namespace()
+    top_down = [top]
+    seen = {top.head.ref}
+    base = top.base.ref
+    while namespace.contains(base):
+        if base in seen:
+            raise CliError(
+                t"Pull request base branches point at each other in a loop, starting at "
+                t"{ui.bookmark(base)}.",
+                hint=t"Retarget one of those pull requests on GitHub so the stack has a "
+                t"bottom, then rerun {ui.cmd('jj-stack checkout')}.",
+            )
+        seen.add(base)
+        try:
+            matches = (await github_client.get_open_prs_by_head_refs(head_refs=(base,))).get(
+                base,
+                (),
+            )
+        except GithubClientError as error:
+            raise CliError(f"Could not inspect pull request branch {base}") from error
+        if len(matches) != 1:
+            raise CliError(
+                t"Expected one pull request for managed base branch {ui.bookmark(base)}, "
+                t"but GitHub reports {len(matches)}.",
+                hint=t"Select the intended PR explicitly and repair it with {ui.cmd('relink')}.",
+            )
+        parent = matches[0]
+        _validate_same_repo_managed_pr(
+            pr=parent,
+            repo=repo,
+        )
+        top_down.append(parent)
+        base = parent.base.ref
+    return tuple(reversed(top_down))
+
+
+def _save_checkout_tracking(
+    *,
+    context: CommandContext,
+    prs: tuple[GithubPR, ...],
+    remote_name: str,
+    repo: GithubRepoAddress,
+    stack: LocalStack,
+    state: TrackingState,
+) -> int:
+    pr_heads = tuple(_require_pr_head_sha(pr) for pr in prs)
+    if len(prs) != len(stack.changes) or pr_heads != tuple(
+        change.commit_id for change in stack.changes
+    ):
+        raise CliError(
+            "The selected pull requests do not describe the stack that was just fetched.",
+            hint=t"Run {ui.cmd('jj-stack view')} to compare them, then submit or "
+            t"relink the pull requests that should match this history.",
+        )
+    remote_targets = context.jj_client.list_remote_branches(
+        remote=remote_name,
+        patterns=tuple(f"refs/heads/{pr.head.ref}" for pr in prs),
+    )
+    replacements: dict[str, tuple[PRIdentity, SubmittedBaseline]] = {}
+    for pr, head_sha, change in zip(
+        prs,
+        pr_heads,
+        stack.changes,
+        strict=True,
+    ):
+        _require_branch_matches_change(branch=pr.head.ref, change=change)
+        if remote_targets.get(pr.head.ref) != head_sha:
+            raise CliError(
+                t"PR #{pr.number} and branch "
+                t"{ui.bookmark(pr.head.ref)} no longer identify the same commit."
+            )
+        replacements[change.change_id] = (
+            PRIdentity(
+                repo_owner=repo.owner,
+                repo_name=repo.repo,
+                pr_number=pr.number,
+                head_owner=repo.owner,
+                head_ref=pr.head.ref,
+            ),
+            SubmittedBaseline(commit_id=head_sha),
+        )
+    _reject_duplicate_checkout_claims(
+        current=state.pr_identities,
+        replacements={change_id: pair[0] for change_id, pair in replacements.items()},
+    )
+    changed_count = sum(
+        (
+            state.pr_identities.get(change_id),
+            state.submitted_baselines.get(change_id),
+        )
+        != replacement
+        for change_id, replacement in replacements.items()
+    )
+    if not changed_count:
+        return 0
+    context.state_store.relink_prs(
+        replacements=replacements,
+    )
+    return changed_count
+
+
+def _reject_duplicate_checkout_claims(
+    *,
+    current: dict[str, PRIdentity],
+    replacements: dict[str, PRIdentity],
+) -> None:
+    combined = dict(current)
+    combined.update(replacements)
+    if duplicate_pr_claim_change_ids(combined).intersection(replacements):
+        raise CliError(
+            "Another saved change already claims one of those pull request numbers or branches.",
+            hint=t"Run {ui.cmd('jj-stack list')} to find the claiming change, then drop its "
+            t"tracking with {ui.cmd('jj-stack unstack --local')} or clean it up with "
+            t"{ui.cmd('jj-stack cleanup')}.",
+        )
+
+
+def _validate_same_repo_managed_pr(
+    *,
+    pr: GithubPR,
+    repo: GithubRepoAddress,
+) -> None:
+    namespace = current_pr_branch_namespace()
+    expected_label = f"{repo.owner}:{pr.head.ref}"
+    if pr.head.label != expected_label:
+        raise CliError(
+            t"Pull request #{pr.number} head "
+            t"{ui.bookmark(pr.head.label or pr.head.ref)} does not "
+            t"belong to {repo.full_name}."
+        )
+    if not namespace.contains(pr.head.ref):
+        raise CliError(
+            t"Pull request #{pr.number} head "
+            t"{ui.bookmark(pr.head.ref)} is not in the reserved "
+            t"{ui.bookmark(namespace.branch_prefix)} namespace."
+        )
+
+
+def _require_pr_head_sha(pr: GithubPR) -> str:
+    head_sha = pr.head.sha
+    if head_sha is None:
+        raise CliError(
+            t"GitHub did not report a head commit for PR #{pr.number}.",
+            hint="Refresh the pull request on GitHub, then retry.",
+        )
+    return head_sha
+
+
+def _require_branch_matches_change(*, branch: str, change: LocalCommit) -> None:
+    if not pr_branch_matches_change(branch, change.change_id):
+        raise CliError(
+            t"PR branch {ui.bookmark(branch)} does not match change "
+            t"{ui.change_id(change.change_id)}."
+        )
+
+
+async def _pick_stack(context: CommandContext) -> CheckoutPickerChoice:
+    """Prompt for one local or GitHub stack without holding the operation lock."""
+
+    state = context.state_store.load()
+    if not state.pr_identities:
+        local_stacks: list[LocalStack] = []
+    else:
+        repo_paths = observe_repo_paths(
+            jj_client=context.jj_client,
+            state=state,
+        )
+        local_stacks = sorted(
+            (path.stack for path in repo_paths.paths if path.tracked_change_ids),
+            key=lambda stack: stack.head.change_id,
+        )
+    remote = select_submit_remote(context.jj_client.list_git_remotes())
+    repo = require_github_repo(remote)
+    async with build_github_client(repo=repo) as github_client:
+        repo_result, stacks_result = await asyncio.gather(
+            github_client.get_repo(),
+            github_client.list_stacks(),
+            return_exceptions=True,
+        )
+        if isinstance(repo_result, GithubClientError):
+            raise CliError(f"Could not inspect GitHub repo {repo.full_name}") from repo_result
+        if isinstance(repo_result, BaseException):
+            raise repo_result
+        if isinstance(stacks_result, GithubClientError):
+            unavailable = github_stacks_unavailable_error(
+                error=stacks_result,
+                repo=repo.full_name,
+            )
+            if unavailable is not None:
+                raise unavailable from None
+            raise CliError("Could not list GitHub stacks for checkout.") from stacks_result
+        if isinstance(stacks_result, BaseException):
+            raise stacks_result
+        github_stacks = stacks_result
+        try:
+            prs = await github_client.get_prs_by_numbers(
+                pr_numbers=tuple(
+                    member.number for stack in github_stacks for member in stack.prs
+                ),
+            )
+        except GithubClientError as error:
+            raise CliError("Could not list GitHub stacks for checkout.") from error
+    choices = _picker_choices(
+        github_stacks=github_stacks,
+        local_stacks=local_stacks,
+        prs=prs,
+        repo=repo,
+        state=state,
+        visible_commit_ids={
+            commit.commit_id
+            for commit in context.jj_client.query_commits_by_ids(
+                tuple(member.head.sha for stack in github_stacks for member in stack.prs)
+            )
+            if not commit.hidden
+        },
+    )
+    return _prompt_picker_choice(choices)
+
+
+def _prompt_picker_choice(
+    choices: tuple[CheckoutPickerChoice, ...],
+) -> CheckoutPickerChoice:
+    """Read one validated numbered selection from the interactive picker."""
+
+    if not choices:
+        raise CliError(
+            "No active local or GitHub stacks to pick from.",
+            hint=t"Use {ui.cmd('checkout --pull-request PR')} to attach a pull request directly.",
+        )
+    console.output("Available stacks:")
+    for index, choice in enumerate(choices, start=1):
+        console.output(f"  [{index}] {choice.heading}")
+        for detail in choice.details:
+            console.output(f"      {detail}")
+    console.output(t"Pick a stack [1-{len(choices)}]: ")
+    selection = sys.stdin.readline().strip()
+    if not selection.isdigit() or not 1 <= int(selection) <= len(choices):
+        raise UsageError(
+            t"{ui.cmd(selection or '(empty)')} is not a valid stack number; "
+            t"expected 1-{len(choices)}."
+        )
+    return choices[int(selection) - 1]
+
+
+def _picker_choices(
+    *,
+    github_stacks: tuple[GithubStack, ...],
+    local_stacks: list[LocalStack],
+    prs: dict[int, GithubPR | None],
+    repo: GithubRepoAddress,
+    state: TrackingState,
+    visible_commit_ids: set[str],
+) -> tuple[CheckoutPickerChoice, ...]:
+    saved_by_pr = {
+        identity.pr_number: (change_id, identity)
+        for change_id, identity in state.pr_identities.items()
+        if identity.repo_key == repo.repo_key
+    }
+    choices: list[CheckoutPickerChoice] = []
+    listed_pr_numbers: set[int] = set()
+    for stack in sorted(github_stacks, key=lambda candidate: candidate.number):
+        active_numbers = stack.active_pr_numbers
+        if not active_numbers:
+            continue
+        numbers = stack.pr_numbers
+        members = tuple(prs.get(number) for number in numbers)
+        if any(member is None for member in members):
+            missing = next(
+                number for number, member in zip(numbers, members, strict=True) if member is None
+            )
+            raise CliError(f"GitHub stack #{stack.number} refers to missing PR #{missing}.")
+        resolved = tuple(member for member in members if member is not None)
+        if not all(_picker_pr_is_adoptable(member, repo) for member in resolved):
+            continue
+        bottom = resolved[0]
+        top = next(member for member in reversed(resolved) if member.number in active_numbers)
+        statuses = Counter(_picker_pr_status(member) for member in resolved)
+        status = ", ".join(
+            f"{count} {name}"
+            for name in ("open", "draft", "closed", "merged")
+            if (count := statuses[name])
+        )
+        active_members = tuple(member for member in resolved if member.number in active_numbers)
+        change_id_by_pr = {
+            member.number: saved[0]
+            for member in active_members
+            if (saved := saved_by_pr.get(member.number)) is not None
+            and saved[1].matches_pr(member)
+        }
+        local = len(change_id_by_pr) == len(active_members) and all(
+            member.head.sha in visible_commit_ids for member in active_members
+        )
+        visible_count = sum(member.head.sha in visible_commit_ids for member in active_members)
+        locality = "local" if local else "partly local" if visible_count else "GitHub only"
+        noun = "PR" if len(numbers) == 1 else "PRs"
+        choices.append(
+            CheckoutPickerChoice(
+                heading=f"GitHub stack #{stack.number} ({locality})",
+                details=(
+                    f"Top: PR #{top.number} {top.title}",
+                    f"Base: {bottom.base.ref}",
+                    f"Size: {len(numbers)} {noun}",
+                    f"Status: {status}",
+                ),
+                revset=change_id_by_pr[top.number] if local else None,
+                pr=None if local else str(top.number),
+            )
+        )
+        listed_pr_numbers.update(numbers)
+    for stack in local_stacks:
+        identity = state.pr_identities.get(stack.head.change_id)
+        if identity is not None and identity.pr_number in listed_pr_numbers:
+            continue
+        count = len(stack.changes)
+        noun = "change" if count == 1 else "changes"
+        choices.append(
+            CheckoutPickerChoice(
+                heading=f"Local stack {stack.head.change_id}",
+                details=(f"Head: {stack.head.subject}", f"Size: {count} {noun}"),
+                revset=stack.head.change_id,
+            )
+        )
+    return tuple(choices)
+
+
+def _picker_pr_status(pr: GithubPR) -> str:
+    normalized = pr.normalize_state()
+    if normalized.state == "open" and normalized.is_draft:
+        return "draft"
+    return normalized.state
+
+
+def _picker_pr_is_adoptable(
+    pr: GithubPR,
+    repo: GithubRepoAddress,
+) -> bool:
+    return current_pr_branch_namespace().contains(pr.head.ref) and (
+        pr.head.label == f"{repo.owner}:{pr.head.ref}"
+    )