erk 0.4.5__py3-none-any.whl

erk/cli/commands/stack/consolidate_cmd.py

@@ -0,0 +1,470 @@
+"""Consolidate worktrees by removing others containing branches from current stack."""
+
+import time
+from pathlib import Path
+
+import click
+
+from erk.cli.activation import render_activation_script
+from erk.cli.commands.branch.unassign_cmd import execute_unassign
+from erk.cli.commands.navigation_helpers import find_assignment_by_worktree_path
+from erk.cli.core import discover_repo_context, worktree_path_for
+from erk.cli.graphite_command import GraphiteCommandWithHiddenOptions
+from erk.cli.help_formatter import script_option
+from erk.core.consolidation_utils import calculate_stack_range, create_consolidation_plan
+from erk.core.context import ErkContext, create_context
+from erk.core.repo_discovery import RepoContext, ensure_erk_metadata_dir
+from erk.core.worktree_pool import load_pool_state
+from erk_shared.git.abc import WorktreeInfo
+from erk_shared.output.output import user_confirm, user_output
+
+
+def _format_section_header(text: str, separator_length: int = 59) -> str:
+    """Format a section header with styled text and separator line."""
+    header = click.style(text, bold=True)
+    separator = "─" * separator_length
+    return f"{header}\n{separator}"
+
+
+def _format_consolidation_plan(
+    stack_branches: list[str],
+    current_branch: str,
+    consolidated_branches: list[str],
+    target_name: str,
+    worktrees_to_remove: list[tuple[str, Path]],
+) -> str:
+    """Format the consolidation plan section with visual hierarchy."""
+    lines: list[str] = []
+
+    # Section header
+    lines.append(_format_section_header("📋 Consolidation Plan"))
+    lines.append("")
+
+    # Branches consolidating to current worktree
+    lines.append("Branches consolidating to current worktree:")
+    for branch in consolidated_branches:
+        if branch == current_branch:
+            branch_display = click.style(branch, fg="bright_green", bold=True)
+            lines.append(f"  • {branch_display} ← (keeping this worktree)")
+        else:
+            lines.append(f"  • {branch}")
+
+    lines.append("")
+
+    # Worktrees to remove
+    lines.append("Worktrees to remove:")
+    for branch, path in worktrees_to_remove:
+        lines.append(f"  • {branch}")
+        lines.append(f"    {click.style(str(path), fg='white', dim=True)}")
+
+    lines.append("")
+    lines.append("─" * 59)
+
+    return "\n".join(lines)
+
+
+def _format_removal_progress(removed_paths: list[Path], unassigned_slots: list[str]) -> str:
+    """Format the removal execution output with grouped checkmarks."""
+    lines: list[str] = []
+
+    if removed_paths or unassigned_slots:
+        lines.append(_format_section_header("🗑️  Removing worktrees..."))
+        for path in removed_paths:
+            lines.append(f"  ✓ {click.style(str(path), fg='green')}")
+        for slot in unassigned_slots:
+            lines.append(f"  ✓ {click.style(slot, fg='cyan')} (slot unassigned)")
+
+    return "\n".join(lines)
+
+
+def _remove_worktree_slot_aware(
+    ctx: ErkContext,
+    repo: RepoContext,
+    wt: WorktreeInfo,
+) -> tuple[Path | None, str | None]:
+    """Remove a worktree with slot awareness.
+
+    If worktree is a pool slot: unassigns slot (keeps directory for reuse).
+    If not a pool slot: removes worktree directory.
+
+    Args:
+        ctx: ErkContext with git operations
+        repo: Repository context
+        wt: WorktreeInfo for the worktree to remove
+
+    Returns:
+        Tuple of (removed_path, unassigned_slot_name):
+        - removed_path: Path if worktree was removed, None if slot unassigned
+        - unassigned_slot_name: Slot name if slot unassigned, None if worktree removed
+    """
+    state = load_pool_state(repo.pool_json_path)
+    assignment = None
+    if state is not None:
+        assignment = find_assignment_by_worktree_path(state, wt.path)
+
+    if assignment is not None:
+        # Slot worktree: unassign instead of remove
+        # state is guaranteed to be non-None since assignment was found in it
+        assert state is not None
+        execute_unassign(ctx, repo, state, assignment)
+        return (None, assignment.slot_name)
+    else:
+        # Non-slot worktree: remove normally
+        ctx.git.remove_worktree(repo.root, wt.path, force=True)
+        return (wt.path, None)
+
+
+@click.command("consolidate", cls=GraphiteCommandWithHiddenOptions)
+@click.argument("branch", required=False, default=None)
+@click.option(
+    "--name",
+    type=str,
+    default=None,
+    help="Create and consolidate into a new worktree with this name",
+)
+@click.option("-f", "--force", is_flag=True, help="Skip confirmation prompt")
+@click.option(
+    "--dry-run",
+    is_flag=True,
+    default=False,
+    help="Show what would be removed without executing",
+)
+@click.option(
+    "--down",
+    is_flag=True,
+    help="Only consolidate downstack (trunk to current branch). Default is entire stack.",
+)
+@script_option
+@click.pass_obj
+def consolidate_stack(
+    ctx: ErkContext,
+    branch: str | None,
+    name: str | None,
+    force: bool,
+    dry_run: bool,
+    down: bool,
+    script: bool,
+) -> None:
+    """Consolidate stack branches into a single worktree.
+
+    By default, consolidates full stack (trunk to leaf). With --down, consolidates
+    only downstack branches (trunk to current).
+
+    This command removes other worktrees that contain branches from the stack,
+    ensuring branches exist in only one worktree. This is useful before
+    stack-wide operations like 'gt restack'.
+
+    BRANCH: Optional branch name. If provided, consolidate only from trunk up to
+    this branch (partial consolidation). Cannot be used with --down.
+
+    \b
+    Examples:
+      # Consolidate full stack into current worktree (default)
+      $ erk consolidate
+
+      # Consolidate only downstack (trunk to current)
+      $ erk consolidate --down
+
+      # Consolidate trunk → feat-2 only (leaves feat-3+ in separate worktrees)
+      $ erk consolidate feat-2
+
+      # Create new worktree "my-stack" and consolidate full stack into it
+      $ erk consolidate --name my-stack
+
+      # Consolidate downstack into new worktree
+      $ erk consolidate --down --name my-partial
+
+      # Preview changes without executing
+      $ erk consolidate --dry-run
+
+      # Skip confirmation prompt
+      $ erk consolidate --force
+
+    Safety checks:
+    - Aborts if any worktree being consolidated has uncommitted changes
+    - Preserves the current worktree (or creates new one with --name)
+    - Shows preview before removal (unless --force)
+    - Never removes root worktree
+    """
+    # During dry-run, always show output regardless of shell integration
+    if dry_run:
+        script = False
+
+    # Validate that --down and BRANCH are not used together
+    if down and branch is not None:
+        user_output(click.style("❌ Error: Cannot use --down with BRANCH argument", fg="red"))
+        user_output(
+            "Use either --down (consolidate trunk to current) or "
+            "BRANCH (consolidate trunk to BRANCH)"
+        )
+        raise SystemExit(1)
+
+    # Get current worktree and branch
+    current_worktree = ctx.cwd
+    current_branch = ctx.git.get_current_branch(current_worktree)
+
+    if current_branch is None:
+        user_output("Error: Current worktree is in detached HEAD state")
+        user_output("Checkout a branch before running consolidate")
+        raise SystemExit(1)
+
+    # Get repository root
+    repo = discover_repo_context(ctx, current_worktree)
+    ensure_erk_metadata_dir(repo)
+
+    # Get current branch's stack
+    stack_branches = ctx.graphite.get_branch_stack(ctx.git, repo.root, current_branch)
+    if stack_branches is None:
+        user_output(f"Error: Branch '{current_branch}' is not tracked by Graphite")
+        user_output(
+            "Run 'gt repo init' to initialize Graphite, or use 'gt track' to track this branch"
+        )
+        raise SystemExit(1)
+
+    # Validate branch argument if provided
+    if branch is not None:
+        if branch not in stack_branches:
+            user_output(
+                click.style(f"❌ Error: Branch '{branch}' is not in the current stack", fg="red")
+            )
+            user_output("\nCurrent stack:")
+            for b in stack_branches:
+                marker = " ← current" if b == current_branch else ""
+                user_output(f"  {click.style(b, fg='cyan')}{marker}")
+            raise SystemExit(1)
+
+    # Determine which portion of the stack to consolidate (now handled by utility)
+    # This will be used in create_consolidation_plan() below
+
+    # Get all worktrees
+    all_worktrees = ctx.git.list_worktrees(repo.root)
+
+    # Validate --name argument if provided
+    if name is not None:
+        # Check if a worktree with this name already exists
+        existing_names = [wt.path.name for wt in all_worktrees]
+
+        if name in existing_names:
+            user_output(click.style(f"❌ Error: Worktree '{name}' already exists", fg="red"))
+            user_output("\nSuggested action:")
+            user_output("  1. Use a different name")
+            user_output(f"  2. Remove existing worktree: erk remove {name}")
+            user_output("  3. Switch to existing: erk br co <branch>")
+            raise SystemExit(1)
+
+    # Calculate stack range early (needed for safety check)
+    # If --down is set, force end_branch to be current_branch
+    end_branch = current_branch if down else branch
+    stack_to_consolidate = calculate_stack_range(stack_branches, end_branch)
+
+    # Check worktrees in stack for uncommitted changes
+    # Only check worktrees that will actually be removed (skip root and current)
+    worktrees_with_changes: list[Path] = []
+    for wt in all_worktrees:
+        if wt.branch not in stack_to_consolidate:
+            continue
+        # Skip root worktree (never removed)
+        if wt.is_root:
+            continue
+        # Skip current worktree (consolidation target, never removed)
+        if wt.path.resolve() == current_worktree.resolve():
+            continue
+        if ctx.git.path_exists(wt.path) and ctx.git.has_uncommitted_changes(wt.path):
+            worktrees_with_changes.append(wt.path)
+
+    if worktrees_with_changes:
+        user_output(
+            click.style("Error: Uncommitted changes detected in worktrees:", fg="red", bold=True)
+        )
+        for wt_path in worktrees_with_changes:
+            user_output(f"  - {wt_path}")
+        user_output("\nCommit or stash changes before running consolidate")
+        raise SystemExit(1)
+
+    # Safety check passed - all worktrees are clean
+    user_output(
+        click.style("✅ Safety check: All worktrees have no uncommitted changes", fg="green")
+    )
+    user_output()
+
+    # Create new worktree if --name is provided
+    # Track temp branch name for cleanup after source worktree removal
+    temp_branch_name: str | None = None
+
+    if name is not None:
+        if not dry_run:
+            # Generate temporary branch name to avoid "already used by worktree" error
+            # when the source worktree and new worktree would have the same branch checked out
+            temp_branch_name = f"temp-consolidate-{int(time.time())}"
+
+            # Use proper erks directory path resolution
+            new_worktree_path = worktree_path_for(repo.worktrees_dir, name)
+
+            # Create temporary branch on current commit (doesn't checkout)
+            # Git operations use check=True, so failures raise CalledProcessError
+            ctx.git.create_branch(current_worktree, temp_branch_name, current_branch)
+
+            # Checkout temporary branch in source worktree to free up the original branch
+            ctx.git.checkout_branch(current_worktree, temp_branch_name)
+
+            # Track temporary branch with Graphite
+            ctx.graphite.track_branch(current_worktree, temp_branch_name, current_branch)
+
+            # Create new worktree with original branch
+            # (now available since source is on temp branch)
+            ctx.git.add_worktree(
+                repo.root,
+                new_worktree_path,
+                branch=current_branch,
+                ref=None,
+                create_branch=False,
+            )
+
+            user_output(click.style(f"✅ Created new worktree: {name}", fg="green"))
+
+            # Change to new worktree directory BEFORE removing source worktree
+            # This prevents the shell from being in a deleted directory
+            # Always change directory regardless of script mode to ensure we're not in
+            # the source worktree when it gets deleted
+            if ctx.git.safe_chdir(new_worktree_path):
+                # Regenerate context with new cwd (context is immutable)
+                ctx = create_context(dry_run=ctx.dry_run)
+                user_output(click.style("✅ Changed directory to new worktree", fg="green"))
+
+            target_worktree_path = new_worktree_path
+        else:
+            user_output(
+                click.style(f"[DRY RUN] Would create new worktree: {name}", fg="yellow", bold=True)
+            )
+            target_worktree_path = current_worktree  # In dry-run, keep current path
+    else:
+        # Use current worktree as target (existing behavior)
+        target_worktree_path = current_worktree
+
+    # Create consolidation plan using utility function
+    # Use the same end_branch logic as calculated above
+    plan = create_consolidation_plan(
+        all_worktrees=all_worktrees,
+        stack_branches=stack_branches,
+        end_branch=end_branch,
+        target_worktree_path=target_worktree_path,
+        source_worktree_path=current_worktree if name is not None else None,
+    )
+
+    # Extract data from plan for easier reference
+    worktrees_to_remove = plan.worktrees_to_remove
+    stack_to_consolidate = plan.stack_to_consolidate
+
+    # Display preview
+    if not worktrees_to_remove:
+        # If using --name, we still need to remove source worktree even if no other worktrees exist
+        if name is None:
+            user_output("No other worktrees found containing branches from current stack")
+            user_output(f"\nCurrent stack branches: {', '.join(stack_branches)}")
+            return
+        # Continue to source worktree removal when using --name
+
+    # Collect data for formatted output
+    worktrees_to_remove_list: list[tuple[str, Path]] = [
+        (wt.branch or "detached", wt.path) for wt in worktrees_to_remove
+    ]
+
+    # Add source worktree to removal list if creating new worktree
+    if name is not None:
+        worktrees_to_remove_list.append((current_branch, current_worktree))
+
+    # Display consolidation plan
+    user_output()
+    plan_output = _format_consolidation_plan(
+        stack_branches=stack_branches,
+        current_branch=current_branch,
+        consolidated_branches=stack_to_consolidate,
+        target_name=name if name is not None else str(current_worktree.name),
+        worktrees_to_remove=worktrees_to_remove_list,
+    )
+    user_output(plan_output)
+
+    # Exit if dry-run
+    if dry_run:
+        user_output(f"\n{click.style('[DRY RUN] No changes made', fg='yellow', bold=True)}")
+        return
+
+    # Get confirmation unless --force or --script
+    if not force and not script:
+        user_output()
+        if not user_confirm("All worktrees are clean. Proceed with removal?", default=False):
+            user_output(click.style("⭕ Aborted", fg="red", bold=True))
+            return
+
+    # Shell integration: generate script to activate new worktree BEFORE destructive operations
+    # This ensures the shell can navigate even if later steps fail (e.g., branch deletion).
+    # The handler will use this script instead of passthrough when available.
+    if name is not None and script and not dry_run:
+        script_content = render_activation_script(
+            worktree_path=target_worktree_path,
+            target_subpath=None,
+            post_cd_commands=None,
+            final_message='echo "✓ Went to consolidated worktree."',
+            comment="work activate-script (consolidate)",
+        )
+        activation_result = ctx.script_writer.write_activation_script(
+            script_content,
+            command_name="consolidate",
+            comment=f"activate {name}",
+        )
+        activation_result.output_for_shell_integration()
+
+    # Remove worktrees and collect paths for progress output
+    removed_paths: list[Path] = []
+    unassigned_slots: list[str] = []
+
+    for wt in worktrees_to_remove:
+        removed_path, slot_name = _remove_worktree_slot_aware(ctx, repo, wt)
+        if removed_path is not None:
+            removed_paths.append(removed_path)
+        if slot_name is not None:
+            unassigned_slots.append(slot_name)
+
+    # Remove source worktree if a new worktree was created
+    if name is not None:
+        # Create a temporary WorktreeInfo for the source worktree
+        source_wt = WorktreeInfo(
+            path=current_worktree.resolve(),
+            branch=current_branch,
+            is_root=False,
+        )
+        removed_path, slot_name = _remove_worktree_slot_aware(ctx, repo, source_wt)
+        if removed_path is not None:
+            removed_paths.append(removed_path)
+        if slot_name is not None:
+            unassigned_slots.append(slot_name)
+
+        # Delete temporary branch after source worktree is removed
+        # (can't delete while it's checked out in the source worktree)
+        if temp_branch_name is not None:
+            ctx.git.delete_branch(repo.root, temp_branch_name, force=True)
+
+    # Display grouped removal progress
+    user_output()
+    user_output(_format_removal_progress(removed_paths, unassigned_slots))
+
+    # Prune stale worktree metadata after all removals
+    # (explicit call now that remove_worktree no longer auto-prunes)
+    ctx.git.prune_worktrees(repo.root)
+
+    user_output(f"\n{click.style('✅ Consolidation complete', fg='green', bold=True)}")
+    user_output()
+    user_output("Next step:")
+    user_output("  Run 'gt restack' to update branch relationships")
+
+    # Early return when no worktree switch (consolidating into current worktree)
+    # Makes it explicit that no script is needed in this case
+    if name is None:
+        return  # No script needed when not switching worktrees
+
+    # Manual cd instruction when not in script mode
+    # (Script mode already output activation script earlier, before destructive operations)
+    if not script and not dry_run:
+        user_output(f"Going to worktree: {click.style(name, fg='cyan', bold=True)}")
+        user_output(f"\n{click.style('ℹ️', fg='blue')} Run this command to switch:")
+        user_output(f"  cd {target_worktree_path}")

erk/cli/commands/stack/list_cmd.py

@@ -0,0 +1,79 @@
+"""List worktree stack with branch info."""
+
+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.cli.graphite import find_worktrees_containing_branch
+from erk.cli.graphite_command import GraphiteCommand
+from erk.core.context import ErkContext
+
+
+@alias("ls")
+@click.command("list", cls=GraphiteCommand)
+@click.pass_obj
+def list_stack(ctx: ErkContext) -> None:
+    """List the worktree stack with branch info.
+
+    Shows branches in the current stack that have associated worktrees,
+    displayed top-to-bottom (upstack children at top, downstack trunk at bottom).
+
+    Table columns:
+    - Marker: → for current branch
+    - branch: Branch name
+    - worktree: Worktree directory name
+    """
+    repo = discover_repo_context(ctx, ctx.cwd)
+    current_branch = ctx.git.get_current_branch(repo.root)
+
+    if current_branch is None:
+        click.echo("Error: Not on a branch (detached HEAD state)", err=True)
+        raise SystemExit(1)
+
+    # Get the stack for current branch
+    stack_branches = ctx.graphite.get_branch_stack(ctx.git, repo.root, current_branch)
+    if stack_branches is None:
+        click.echo(f"Error: Branch '{current_branch}' is not tracked by Graphite", err=True)
+        click.echo("Run 'gt track' to track this branch, or 'gt create' to create a new branch.")
+        raise SystemExit(1)
+
+    # Get worktrees for branch-to-worktree mapping
+    worktrees = ctx.git.list_worktrees(repo.root)
+
+    # Build table
+    table = Table(show_header=True, header_style="bold", box=None)
+    table.add_column("", no_wrap=True)  # Marker column
+    table.add_column("branch", style="cyan", no_wrap=True)
+    table.add_column("worktree", no_wrap=True)
+
+    # Show all branches in stack, using ancestor worktree for branches without their own
+    # Reverse order: upstack (children) at top, downstack (trunk) at bottom
+    for branch in reversed(stack_branches):
+        matching_worktrees = find_worktrees_containing_branch(ctx, repo.root, worktrees, branch)
+
+        if matching_worktrees:
+            wt = matching_worktrees[0]
+        else:
+            # Branch has no direct worktree - find closest ancestor with one
+            wt = ctx.graphite.find_ancestor_worktree(ctx.git, repo.root, branch)
+            if wt is None:
+                continue  # Only skip if truly no worktree found
+
+        wt_name = "root" if wt.is_root else wt.path.name
+
+        is_current = branch == current_branch
+        marker = "→" if is_current else ""
+        branch_display = f"[bold cyan]{branch}[/bold cyan]" if is_current else branch
+
+        table.add_row(marker, branch_display, wt_name)
+
+    # Check if table has any rows
+    if table.row_count == 0:
+        click.echo("No branches in stack have worktrees", err=True)
+        return
+
+    # Output table to stderr (consistent with erk wt list)
+    console = Console(stderr=True, force_terminal=True)
+    console.print(table)