erk 0.4.5__py3-none-any.whl

erk/cli/commands/branch/checkout_cmd.py

@@ -0,0 +1,357 @@
+"""Checkout command - find and switch to a worktree by branch name."""
+
+from pathlib import Path
+
+import click
+
+from erk.cli.activation import render_activation_script
+from erk.cli.alias import alias
+from erk.cli.commands.completions import complete_branch_names
+from erk.cli.commands.wt.create_cmd import ensure_worktree_for_branch
+from erk.cli.core import discover_repo_context
+from erk.cli.graphite import find_worktrees_containing_branch
+from erk.cli.help_formatter import CommandWithHiddenOptions, script_option
+from erk.core.context import ErkContext
+from erk.core.repo_discovery import RepoContext, ensure_erk_metadata_dir
+from erk.core.worktree_utils import compute_relative_path_in_worktree
+from erk_shared.git.abc import WorktreeInfo
+from erk_shared.output.output import user_confirm, user_output
+
+
+def try_switch_root_worktree(ctx: ErkContext, repo: RepoContext, branch: str) -> Path | None:
+    """Try to switch root worktree to branch if it's trunk and root is clean.
+
+    This implements the "takeover" behavior where checking out trunk in a clean root
+    worktree switches the root to trunk instead of creating a new dated worktree.
+
+    Args:
+        ctx: Erk context with git operations
+        repo: Repository context
+        branch: Branch name to check
+
+    Returns:
+        Root worktree path if successful, None otherwise
+    """
+    # Check if branch is trunk
+    if branch != ctx.trunk_branch:
+        return None
+
+    # Find root worktree
+    worktrees = ctx.git.list_worktrees(repo.root)
+    root_worktree = None
+    for wt in worktrees:
+        if wt.is_root:
+            root_worktree = wt
+            break
+
+    if root_worktree is None:
+        return None
+
+    # Check if root is clean
+    if not ctx.git.is_worktree_clean(root_worktree.path):
+        return None
+
+    # Switch root to trunk branch
+    ctx.git.checkout_branch(root_worktree.path, branch)
+
+    return root_worktree.path
+
+
+def _ensure_graphite_tracking(
+    ctx: ErkContext,
+    repo_root: Path,
+    target_path: Path,
+    branch: str,
+    script: bool,
+) -> None:
+    """Ensure branch is tracked by Graphite (idempotent), with user confirmation.
+
+    If the branch is not already tracked, prompts the user and tracks it with
+    trunk as parent if confirmed. This enables branches created without Graphite
+    (e.g., via erk-queue) to be managed with Graphite locally.
+
+    Args:
+        ctx: Erk context
+        repo_root: Repository root path
+        target_path: Worktree path where `gt track` should run
+        branch: Target branch name
+        script: Whether to output only the activation script
+    """
+    # Skip if Graphite is disabled
+    use_graphite = ctx.global_config.use_graphite if ctx.global_config else False
+    if not use_graphite:
+        return
+
+    trunk_branch = ctx.trunk_branch
+    # Skip if no trunk branch detected (shouldn't happen in checkout context)
+    if trunk_branch is None:
+        return
+
+    # Skip trunk branch - it's always implicitly tracked
+    if branch == trunk_branch:
+        return
+
+    # Check if already tracked (LBYL)
+    all_branches = ctx.graphite.get_all_branches(ctx.git, repo_root)
+    if branch in all_branches:
+        return  # Already tracked, nothing to do
+
+    # In script mode, skip tracking (no interactive prompts allowed)
+    if script:
+        return
+
+    # Prompt user for confirmation
+    if not user_confirm(
+        f"Branch '{branch}' is not tracked by Graphite. Track it with parent '{trunk_branch}'?",
+        default=False,
+    ):
+        return
+
+    # Track the branch with trunk as parent
+    ctx.graphite.track_branch(target_path, branch, trunk_branch)
+    user_output(f"Tracked '{branch}' with Graphite (parent: {trunk_branch})")
+
+
+def _format_worktree_info(wt: WorktreeInfo, repo_root: Path) -> str:
+    """Format worktree information for display.
+
+    Args:
+        wt: WorktreeInfo to format
+        repo_root: Path to repository root (used to identify root worktree)
+
+    Returns:
+        Formatted string like "root (currently on 'main')" or "wt-name (currently on 'feature')"
+    """
+    current = wt.branch or "(detached HEAD)"
+    if wt.path == repo_root:
+        return f"  - root (currently on '{current}')"
+    else:
+        # Get worktree name from path
+        wt_name = wt.path.name
+        return f"  - {wt_name} (currently on '{current}')"
+
+
+def _perform_checkout(
+    ctx: ErkContext,
+    *,
+    repo_root: Path,
+    target_worktree: WorktreeInfo,
+    branch: str,
+    script: bool,
+    is_newly_created: bool = False,
+    worktrees: list[WorktreeInfo] | None = None,
+) -> None:
+    """Perform the actual checkout and switch to a worktree.
+
+    Args:
+        ctx: Erk context
+        repo_root: Repository root path
+        target_worktree: The worktree to switch to
+        branch: Target branch name
+        script: Whether to output only the activation script
+        is_newly_created: Whether the worktree was just created (default False)
+        worktrees: Optional list of worktrees (for relative path computation)
+    """
+    target_path = target_worktree.path
+    current_branch_in_worktree = target_worktree.branch
+    current_cwd = ctx.cwd
+
+    # Compute relative path to preserve directory position
+    relative_path = compute_relative_path_in_worktree(worktrees, ctx.cwd) if worktrees else None
+
+    # Check if branch is already checked out in the worktree
+    need_checkout = current_branch_in_worktree != branch
+
+    # If we need to checkout, do it before generating the activation script
+    if need_checkout:
+        # Checkout the branch in the target worktree
+        ctx.git.checkout_branch(target_path, branch)
+
+    # Ensure branch is tracked with Graphite (idempotent)
+    _ensure_graphite_tracking(ctx, repo_root, target_path, branch, script)
+
+    if need_checkout:
+        # Show stack context
+        if not script:
+            stack = ctx.graphite.get_branch_stack(ctx.git, repo_root, branch)
+            if stack:
+                user_output(f"Stack: {' -> '.join(stack)}")
+            user_output(f"Checked out '{branch}' in worktree")
+
+    # Generate activation script
+    if script:
+        # Script mode: always generate script (for shell integration or manual sourcing)
+        is_switching_location = current_cwd != target_path
+
+        # Determine worktree name from path
+        worktree_name = target_path.name
+
+        # Four-case message logic:
+        if is_newly_created:
+            # Case 4: Switched to newly created worktree
+            styled_wt = click.style(worktree_name, fg="cyan", bold=True)
+            switch_message = f'echo "Switched to new worktree {styled_wt}"'
+        elif not is_switching_location:
+            # Case 1: Already on target branch in current worktree
+            styled_branch = click.style(branch, fg="yellow")
+            styled_wt = click.style(worktree_name, fg="cyan", bold=True)
+            switch_message = f'echo "Already on branch {styled_branch} in worktree {styled_wt}"'
+        elif not need_checkout:
+            # Case 2: Switched to existing worktree with branch already checked out
+            styled_wt = click.style(worktree_name, fg="cyan", bold=True)
+            if worktree_name == branch:
+                # Standard naming
+                switch_message = f'echo "Switched to worktree {styled_wt}"'
+            else:
+                # Edge case: non-standard naming
+                styled_branch = click.style(branch, fg="yellow")
+                switch_message = f'echo "Switched to worktree {styled_wt} (branch {styled_branch})"'
+        else:
+            # Case 3: Switched to existing worktree and checked out branch
+            styled_wt = click.style(worktree_name, fg="cyan", bold=True)
+            styled_branch = click.style(branch, fg="yellow")
+            switch_message = (
+                f'echo "Switched to worktree {styled_wt} and checked out branch {styled_branch}"'
+            )
+
+        script_content = render_activation_script(
+            worktree_path=target_path,
+            target_subpath=relative_path,
+            post_cd_commands=None,
+            final_message=switch_message,
+            comment="work activate-script",
+        )
+
+        result = ctx.script_writer.write_activation_script(
+            script_content,
+            command_name="checkout",
+            comment=f"checkout {branch}",
+        )
+        result.output_for_shell_integration()
+    else:
+        # Non-script mode: Apply same four-case logic with user_output()
+        worktree_name = target_path.name
+
+        if is_newly_created:
+            styled_wt = click.style(worktree_name, fg="cyan", bold=True)
+            user_output(f"Switched to new worktree {styled_wt}")
+        elif ctx.cwd == target_path:
+            styled_branch = click.style(branch, fg="yellow")
+            styled_wt = click.style(worktree_name, fg="cyan", bold=True)
+            user_output(f"Already on branch {styled_branch} in worktree {styled_wt}")
+        elif current_branch_in_worktree == branch:
+            styled_wt = click.style(worktree_name, fg="cyan", bold=True)
+            if worktree_name == branch:
+                user_output(f"Switched to worktree {styled_wt}")
+            else:
+                styled_branch = click.style(branch, fg="yellow")
+                user_output(f"Switched to worktree {styled_wt} (branch {styled_branch})")
+        else:
+            styled_wt = click.style(worktree_name, fg="cyan", bold=True)
+            styled_branch = click.style(branch, fg="yellow")
+            user_output(f"Switched to worktree {styled_wt} and checked out branch {styled_branch}")
+
+        # Show manual instructions
+        user_output("\nShell integration not detected. Run 'erk init --shell' to set up.")
+        user_output(f"Or use: source <(erk br co {branch} --script)")
+
+
+@alias("co")
+@click.command("checkout", cls=CommandWithHiddenOptions)
+@click.argument("branch", metavar="BRANCH", shell_complete=complete_branch_names)
+@script_option
+@click.pass_obj
+def branch_checkout(ctx: ErkContext, branch: str, script: bool) -> None:
+    """Checkout BRANCH by finding and switching to its worktree.
+
+    This command finds which worktree has the specified branch checked out
+    and switches to it. If the branch exists but isn't checked out anywhere,
+    a worktree is automatically created. If the branch exists on origin but
+    not locally, a tracking branch and worktree are created automatically.
+
+    Examples:
+
+        erk br co feature/user-auth      # Checkout existing worktree
+
+        erk br co unchecked-branch       # Auto-create worktree
+
+        erk br co origin-only-branch     # Create tracking branch + worktree
+
+    If multiple worktrees contain the branch, all options are shown.
+    """
+    # Use existing repo from context if available (for tests), otherwise discover
+    if isinstance(ctx.repo, RepoContext):
+        repo = ctx.repo
+    else:
+        repo = discover_repo_context(ctx, ctx.cwd)
+    ensure_erk_metadata_dir(repo)
+
+    # Get all worktrees
+    worktrees = ctx.git.list_worktrees(repo.root)
+
+    # Find worktrees containing the target branch
+    matching_worktrees = find_worktrees_containing_branch(ctx, repo.root, worktrees, branch)
+
+    # Track whether we're creating a new worktree
+    is_newly_created = False
+
+    # Handle three cases: no match, one match, multiple matches
+    if len(matching_worktrees) == 0:
+        # No worktrees have this branch checked out
+        # First, try switching clean root worktree if checking out trunk
+        root_path = try_switch_root_worktree(ctx, repo, branch)
+        if root_path is not None:
+            # Successfully switched root to trunk - refresh and jump to it
+            worktrees = ctx.git.list_worktrees(repo.root)
+            matching_worktrees = find_worktrees_containing_branch(ctx, repo.root, worktrees, branch)
+        else:
+            # Root not available or not trunk - auto-create worktree
+            _worktree_path, is_newly_created = ensure_worktree_for_branch(
+                ctx, repo, branch, is_plan_derived=False
+            )
+
+            # Refresh worktree list to include the newly created worktree
+            worktrees = ctx.git.list_worktrees(repo.root)
+            matching_worktrees = find_worktrees_containing_branch(ctx, repo.root, worktrees, branch)
+
+        # Fall through to jump to the worktree
+
+    if len(matching_worktrees) == 1:
+        # Exactly one worktree contains this branch
+        target_worktree = matching_worktrees[0]
+        _perform_checkout(
+            ctx,
+            repo_root=repo.root,
+            target_worktree=target_worktree,
+            branch=branch,
+            script=script,
+            is_newly_created=is_newly_created,
+            worktrees=worktrees,
+        )
+
+    else:
+        # Multiple worktrees contain this branch
+        # Check if any worktree has the branch directly checked out
+        directly_checked_out = [wt for wt in matching_worktrees if wt.branch == branch]
+
+        if len(directly_checked_out) == 1:
+            # Exactly one worktree has the branch directly checked out - jump to it
+            target_worktree = directly_checked_out[0]
+            _perform_checkout(
+                ctx,
+                repo_root=repo.root,
+                target_worktree=target_worktree,
+                branch=branch,
+                script=script,
+                is_newly_created=is_newly_created,
+                worktrees=worktrees,
+            )
+        else:
+            # Zero or multiple worktrees have it directly checked out
+            # Show error message listing all options
+            user_output(f"Branch '{branch}' exists in multiple worktrees:")
+            for wt in matching_worktrees:
+                user_output(_format_worktree_info(wt, repo.root))
+
+            user_output("\nPlease specify which worktree to use.")
+            raise SystemExit(1)

erk/cli/commands/branch/create_cmd.py

@@ -0,0 +1,161 @@
+"""Branch create command - create a new branch with optional slot assignment."""
+
+import sys
+from datetime import UTC, datetime
+
+import click
+
+from erk.cli.commands.slot.common import (
+    cleanup_worktree_artifacts,
+    find_branch_assignment,
+    find_inactive_slot,
+    find_next_available_slot,
+    generate_slot_name,
+    get_pool_size,
+    handle_pool_full_interactive,
+)
+from erk.cli.core import discover_repo_context
+from erk.core.context import ErkContext
+from erk.core.repo_discovery import ensure_erk_metadata_dir
+from erk.core.worktree_pool import (
+    PoolState,
+    SlotAssignment,
+    load_pool_state,
+    save_pool_state,
+)
+from erk_shared.output.output import user_output
+
+
+@click.command("create")
+@click.argument("branch_name", metavar="BRANCH")
+@click.option("--no-slot", is_flag=True, help="Create branch without slot assignment")
+@click.option("-f", "--force", is_flag=True, help="Auto-unassign oldest branch if pool is full")
+@click.pass_obj
+def branch_create(ctx: ErkContext, branch_name: str, no_slot: bool, force: bool) -> None:
+    """Create a NEW branch and optionally assign it to a pool slot.
+
+    BRANCH is the name of the new git branch to create.
+
+    By default, the command will:
+    1. Verify the branch does NOT already exist (fails if it does)
+    2. Create the branch from trunk
+    3. Find the next available slot in the pool
+    4. Create a worktree for that slot
+    5. Assign the branch to the slot
+
+    Use --no-slot to create a branch without assigning it to a slot.
+    Use `erk br assign` to assign an EXISTING branch to a slot.
+    """
+    repo = discover_repo_context(ctx, ctx.cwd)
+    ensure_erk_metadata_dir(repo)
+
+    # Check if branch already exists
+    local_branches = ctx.git.list_local_branches(repo.root)
+    if branch_name in local_branches:
+        user_output(
+            f"Error: Branch '{branch_name}' already exists.\n"
+            "Use `erk br assign` to assign an existing branch to a slot."
+        )
+        raise SystemExit(1) from None
+
+    # Create the new branch from trunk
+    trunk = ctx.git.detect_trunk_branch(repo.root)
+    ctx.git.create_branch(repo.root, branch_name, trunk)
+    ctx.graphite.track_branch(repo.root, branch_name, trunk)
+    user_output(f"Created branch: {branch_name}")
+
+    # If --no-slot is specified, we're done
+    if no_slot:
+        return
+
+    # Get pool size from config or default
+    pool_size = get_pool_size(ctx)
+
+    # Load or create pool state
+    state = load_pool_state(repo.pool_json_path)
+    if state is None:
+        state = PoolState(
+            version="1.0",
+            pool_size=pool_size,
+            slots=(),
+            assignments=(),
+        )
+
+    # Check if branch is already assigned (shouldn't happen since we just created it)
+    existing = find_branch_assignment(state, branch_name)
+    if existing is not None:
+        user_output(f"Error: Branch '{branch_name}' already assigned to {existing.slot_name}")
+        raise SystemExit(1) from None
+
+    # First, prefer reusing existing worktrees (fast path)
+    inactive_slot = find_inactive_slot(state, ctx.git, repo.root)
+    if inactive_slot is not None:
+        slot_name, worktree_path = inactive_slot
+
+        # Checkout the branch in the existing worktree
+        ctx.git.checkout_branch(worktree_path, branch_name)
+    else:
+        # Fall back to on-demand slot creation
+        slot_num = find_next_available_slot(state, repo.worktrees_dir)
+        if slot_num is None:
+            # Pool is full - handle interactively or with --force
+            to_unassign = handle_pool_full_interactive(state, force, sys.stdin.isatty())
+            if to_unassign is None:
+                raise SystemExit(1) from None
+
+            # Remove the assignment from state
+            new_assignments = tuple(
+                a for a in state.assignments if a.slot_name != to_unassign.slot_name
+            )
+            state = PoolState(
+                version=state.version,
+                pool_size=state.pool_size,
+                slots=state.slots,
+                assignments=new_assignments,
+            )
+            save_pool_state(repo.pool_json_path, state)
+            user_output(
+                click.style("✓ ", fg="green")
+                + f"Unassigned {click.style(to_unassign.branch_name, fg='yellow')} "
+                + f"from {click.style(to_unassign.slot_name, fg='cyan')}"
+            )
+
+            # Reuse the unassigned slot - worktree exists, just checkout
+            slot_name = to_unassign.slot_name
+            worktree_path = to_unassign.worktree_path
+            cleanup_worktree_artifacts(worktree_path)
+            ctx.git.checkout_branch(worktree_path, branch_name)
+        else:
+            # Create new slot - no worktree exists yet
+            slot_name = generate_slot_name(slot_num)
+            worktree_path = repo.worktrees_dir / slot_name
+            worktree_path.mkdir(parents=True, exist_ok=True)
+            ctx.git.add_worktree(
+                repo.root,
+                worktree_path,
+                branch=branch_name,
+                ref=None,
+                create_branch=False,
+            )
+
+    # Create new assignment
+    now = datetime.now(UTC).isoformat()
+    new_assignment = SlotAssignment(
+        slot_name=slot_name,
+        branch_name=branch_name,
+        assigned_at=now,
+        worktree_path=worktree_path,
+    )
+
+    # Update state with new assignment
+    new_state = PoolState(
+        version=state.version,
+        pool_size=state.pool_size,
+        slots=state.slots,
+        assignments=(*state.assignments, new_assignment),
+    )
+
+    # Save state
+    save_pool_state(repo.pool_json_path, new_state)
+
+    user_output(click.style(f"✓ Assigned {branch_name} to {slot_name}", fg="green"))

erk/cli/commands/branch/list_cmd.py

@@ -0,0 +1,82 @@
+"""List active branches - those with worktrees or open PRs."""
+
+import click
+from rich.console import Console
+from rich.table import Table
+
+from erk.cli.alias import alias
+from erk.cli.core import discover_repo_context
+from erk.core.context import ErkContext
+from erk.core.display_utils import format_relative_time, get_pr_status_emoji
+from erk_shared.github.types import PullRequestInfo
+
+
+@alias("ls")
+@click.command("list")
+@click.pass_obj
+def branch_list(ctx: ErkContext) -> None:
+    """List active branches with their worktrees and PR status.
+
+    Active branches are those that have:
+    - A worktree checked out, OR
+    - An open pull request
+
+    Example:
+        erk br ls
+    """
+    repo = discover_repo_context(ctx, ctx.cwd)
+
+    # Detect trunk branch for last commit calculation
+    trunk = ctx.git.detect_trunk_branch(repo.root)
+
+    # Get worktrees and PR info
+    worktrees = ctx.git.list_worktrees(repo.root)
+    prs = ctx.graphite.get_prs_from_graphite(ctx.git, repo.root)
+
+    # Build active branches map: branch -> (worktree_name, pr_info)
+    active_branches: dict[str, tuple[str | None, PullRequestInfo | None]] = {}
+
+    # Add branches from worktrees
+    for wt in worktrees:
+        if wt.branch is not None and wt.branch != ctx.trunk_branch:
+            wt_name = "root" if wt.is_root else wt.path.name
+            pr = prs.get(wt.branch)
+            active_branches[wt.branch] = (wt_name, pr)
+
+    # Add branches with open PRs (not already in worktrees)
+    for branch, pr in prs.items():
+        if branch not in active_branches and pr.state == "OPEN":
+            active_branches[branch] = (None, pr)
+
+    # Display table
+    console = Console(stderr=True, force_terminal=True)
+    table = Table(show_header=True, header_style="bold", box=None)
+    table.add_column("branch", style="yellow")
+    table.add_column("worktree", style="cyan")
+    table.add_column("pr")
+    table.add_column("last", no_wrap=True)
+    table.add_column("state", style="dim")
+
+    for branch in sorted(active_branches.keys()):
+        wt_name, pr = active_branches[branch]
+
+        pr_cell = "-"
+        state_cell = "-"
+        if pr is not None:
+            emoji = get_pr_status_emoji(pr)
+            pr_cell = f"{emoji} #{pr.number}"
+            state_cell = pr.state
+
+        # Get last commit time for this branch
+        timestamp = ctx.git.get_branch_last_commit_time(repo.root, branch, trunk)
+        last_cell = format_relative_time(timestamp) if timestamp is not None else "-"
+
+        table.add_row(
+            branch,
+            wt_name or "-",
+            pr_cell,
+            last_cell,
+            state_cell,
+        )
+
+    console.print(table)