Coverage for src/workstack/cli/commands/jump.py: 83%

1"""Jump command - find and switch to a worktree by branch name.""" 

2 

3import shlex 

4from pathlib import Path 

5 

6import click 

7 

8from workstack.cli.activation import render_activation_script 

9from workstack.cli.core import discover_repo_context 

10from workstack.cli.graphite import find_worktrees_containing_branch, get_branch_stack 

11from workstack.cli.shell_utils import write_script_to_temp 

12from workstack.core.context import WorkstackContext 

13from workstack.core.gitops import WorktreeInfo 

14 

15 

16def _format_worktree_info(wt: WorktreeInfo, repo_root: Path) -> str: 

17 """Format worktree information for display. 

18 

19 Args: 

20 wt: WorktreeInfo to format 

21 repo_root: Path to repository root (used to identify root worktree) 

22 

23 Returns: 

24 Formatted string like "root (currently on 'main')" or "wt-name (currently on 'feature')" 

25 """ 

26 current = wt.branch or "(detached HEAD)" 

27 if wt.path == repo_root: 

28 return f" - root (currently on '{current}')" 

29 else: 

30 # Get worktree name from path 

31 wt_name = wt.path.name 

32 return f" - {wt_name} (currently on '{current}')" 

33 

34 

35def _perform_jump( 

36 ctx: WorkstackContext, 

37 repo_root: Path, 

38 target_worktree: WorktreeInfo, 

39 branch: str, 

40 script: bool, 

41) -> None: 

42 """Perform the actual jump to a worktree. 

43 

44 Args: 

45 ctx: Workstack context 

46 repo_root: Repository root path 

47 target_worktree: The worktree to jump to 

48 branch: Target branch name 

49 script: Whether to output only the activation script 

50 """ 

51 target_path = target_worktree.path 

52 current_branch_in_worktree = target_worktree.branch 

53 

54 # Defensive check: verify path exists before any operations 

55 if not target_path.exists(): 

56 click.echo( 

57 f"Error: Worktree path does not exist: {target_path}", 

58 err=True, 

59 ) 

60 raise SystemExit(1) 

61 

62 # Check if we're already on the target branch in the target worktree 

63 current_cwd = Path.cwd() 

64 if current_cwd == target_path and current_branch_in_worktree == branch: 

65 if not script: 

66 click.echo(f"Already on branch '{branch}' in this worktree") 

67 return 

68 

69 # Check if branch is already checked out in the worktree 

70 need_checkout = current_branch_in_worktree != branch 

71 

72 # If we need to checkout, do it before generating the activation script 

73 if need_checkout: 

74 # Checkout the branch in the target worktree 

75 ctx.git_ops.checkout_branch(target_path, branch) 

76 

77 # Show stack context 

78 if not script: 

79 stack = get_branch_stack(ctx, repo_root, branch) 

80 if stack: 

81 click.echo(f"Stack: {' -> '.join(stack)}") 

82 click.echo(f"Checked out '{branch}' in worktree") 

83 

84 # Generate activation script 

85 if script: 

86 # Script mode: always generate script (for shell integration or manual sourcing) 

87 # Use shlex.quote() for branch name security (defense-in-depth) 

88 safe_branch = shlex.quote(branch) 

89 if need_checkout: 

90 jump_message = f'echo "Jumped to branch {safe_branch}: $(pwd)"' 

91 else: 

92 jump_message = f'echo "Already on branch {safe_branch}: $(pwd)"' 

93 script_content = render_activation_script( 

94 worktree_path=target_path, final_message=jump_message 

95 ) 

96 

97 script_path = write_script_to_temp( 

98 script_content, 

99 command_name="jump", 

100 comment=f"jump to {branch}", 

101 ) 

102 click.echo(str(script_path), nl=False) 

103 else: 

104 # No shell integration available, show manual instructions 

105 click.echo( 

106 "Shell integration not detected. " 

107 "Run 'workstack init --shell' to set up automatic activation." 

108 ) 

109 click.echo(f"\nOr use: source <(workstack jump {branch} --script)") 

110 

111 

112@click.command("jump") 

113@click.argument("branch", metavar="BRANCH") 

114@click.option( 

115 "--script", is_flag=True, help="Print only the activation script without usage instructions." 

116) 

117@click.pass_obj 

118def jump_cmd(ctx: WorkstackContext, branch: str, script: bool) -> None: 

119 """Jump to BRANCH by finding and switching to its worktree. 

120 

121 This command finds which worktree contains the specified branch 

122 in its Graphite stack and switches to it. The branch does not 

123 need to be currently checked out - it just needs to exist in 

124 the worktree's stack lineage. 

125 

126 Requires Graphite to be enabled. 

127 

128 Examples: 

129 

130 workstack jump feature/user-auth 

131 

132 workstack jump hotfix/critical-bug 

133 

134 If multiple worktrees contain the branch, all options are shown. 

135 """ 

136 # Check if Graphite is enabled 

137 if not ctx.global_config_ops.get_use_graphite(): 

138 click.echo( 

139 "Error: Jump command requires Graphite. Enable Graphite with:\n" 

140 " workstack config set use_graphite true", 

141 err=True, 

142 ) 

143 raise SystemExit(1) 

144 

145 repo = discover_repo_context(ctx, Path.cwd()) 

146 

147 # Get all worktrees 

148 worktrees = ctx.git_ops.list_worktrees(repo.root) 

149 

150 # Find worktrees containing the target branch 

151 matching_worktrees = find_worktrees_containing_branch(ctx, repo.root, worktrees, branch) 

152 

153 # Handle three cases: no match, one match, multiple matches 

154 if len(matching_worktrees) == 0: 

155 # No worktrees contain this branch 

156 click.echo( 

157 f"Error: Branch '{branch}' not found in any worktree stack.\n" 

158 f"To create a worktree with this branch, run:\n" 

159 f" workstack create --from-branch {branch}", 

160 err=True, 

161 ) 

162 raise SystemExit(1) 

163 

164 if len(matching_worktrees) == 1: 

165 # Exactly one worktree contains this branch 

166 target_worktree = matching_worktrees[0] 

167 _perform_jump(ctx, repo.root, target_worktree, branch, script) 

168 

169 else: 

170 # Multiple worktrees contain this branch 

171 # Check if any worktree has the branch directly checked out 

172 directly_checked_out = [wt for wt in matching_worktrees if wt.branch == branch] 

173 

174 if len(directly_checked_out) == 1: 

175 # Exactly one worktree has the branch directly checked out - jump to it 

176 target_worktree = directly_checked_out[0] 

177 _perform_jump(ctx, repo.root, target_worktree, branch, script) 

178 else: 

179 # Zero or multiple worktrees have it directly checked out 

180 # Show error message listing all options 

181 click.echo(f"Branch '{branch}' exists in multiple worktrees:", err=True) 

182 for wt in matching_worktrees: 

183 click.echo(_format_worktree_info(wt, repo.root), err=True) 

184 

185 click.echo("\nUse 'workstack switch' to choose a specific worktree first.", err=True) 

186 raise SystemExit(1)