Coverage for src/workstack/cli/tree.py: 99%
121 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-10-19 09:31 -0400
1"""Tree visualization for workstack.
3This module builds and renders tree structures showing worktrees and their
4Graphite dependency relationships. All functions are pure (no I/O) except
5for the entry point which loads data via WorkstackContext.
6"""
8import json
9from dataclasses import dataclass
10from pathlib import Path
12import click
14from workstack.core.context import WorkstackContext
17@dataclass(frozen=True)
18class TreeNode:
19 """A node in the workstack tree.
21 Represents a branch that has an active worktree, with its children
22 (dependent branches that also have worktrees).
24 Attributes:
25 branch_name: Git branch name (e.g., "fix-workstack-s")
26 worktree_name: Worktree directory name (e.g., "root", "fix-plan")
27 children: List of child TreeNode objects
28 is_current: True if this worktree is the current working directory
29 """
31 branch_name: str
32 worktree_name: str
33 children: list["TreeNode"]
34 is_current: bool
37@dataclass(frozen=True)
38class WorktreeMapping:
39 """Mapping between branches and their worktrees.
41 Attributes:
42 branch_to_worktree: Map of branch name -> worktree name
43 worktree_to_path: Map of worktree name -> filesystem path
44 current_worktree: Name of current worktree (None if not in a worktree)
45 """
47 branch_to_worktree: dict[str, str]
48 worktree_to_path: dict[str, Path]
49 current_worktree: str | None
52@dataclass(frozen=True)
53class BranchGraph:
54 """Graph of branch relationships from Graphite cache.
56 Attributes:
57 parent_of: Map of branch name -> parent branch name
58 children_of: Map of branch name -> list of child branch names
59 trunk_branches: List of trunk branch names (branches with no parent)
60 """
62 parent_of: dict[str, str]
63 children_of: dict[str, list[str]]
64 trunk_branches: list[str]
67def build_workstack_tree(
68 ctx: WorkstackContext,
69 repo_root: Path,
70) -> list[TreeNode]:
71 """Build tree structure of ONLY branches with active worktrees.
73 This is the main entry point that orchestrates the tree building process:
74 1. Get all worktrees and their branches from git
75 2. Load Graphite cache for parent-child relationships (REQUIRED)
76 3. Build branch graph from cache data
77 4. Filter graph to ONLY branches that have worktrees
78 5. Build tree starting from trunk branches
79 6. Return list of root nodes (typically just "main")
81 Args:
82 ctx: Workstack context with git operations
83 repo_root: Path to repository root
85 Returns:
86 List of root TreeNode objects (typically one for trunk)
88 Raises:
89 SystemExit: If Graphite cache doesn't exist or can't be loaded
90 """
91 # Step 1: Get worktrees
92 worktree_mapping = _get_worktree_mapping(ctx, repo_root)
94 # Step 2: Load Graphite cache (REQUIRED - hard fail if missing)
95 branch_graph = _load_graphite_branch_graph(ctx, repo_root)
96 if branch_graph is None:
97 click.echo(
98 "Error: Graphite cache not found. The 'tree' command requires Graphite.\n"
99 "Make sure Graphite is enabled: workstack config set use-graphite true",
100 err=True,
101 )
102 raise SystemExit(1)
104 # Step 3: Filter graph to only branches with worktrees
105 active_branches = set(worktree_mapping.branch_to_worktree.keys())
106 filtered_graph = _filter_graph_to_active_branches(branch_graph, active_branches)
108 # Step 4: Build tree from filtered graph
109 return _build_tree_from_graph(filtered_graph, worktree_mapping)
112def _get_worktree_mapping(
113 ctx: WorkstackContext,
114 repo_root: Path,
115) -> WorktreeMapping:
116 """Get mapping of branches to worktrees.
118 Queries git for all worktrees and creates mappings between branches,
119 worktree names, and filesystem paths. Detects the current worktree.
121 Args:
122 ctx: Workstack context with git operations
123 repo_root: Path to repository root
125 Returns:
126 WorktreeMapping with all active worktrees and their branches
127 """
128 worktrees = ctx.git_ops.list_worktrees(repo_root)
129 current_path = Path.cwd().resolve()
131 branch_to_worktree: dict[str, str] = {}
132 worktree_to_path: dict[str, Path] = {}
133 current_worktree: str | None = None
135 for wt in worktrees:
136 # Skip worktrees with detached HEAD
137 if wt.branch is None:
138 continue
140 # Determine worktree name
141 if wt.path.resolve() == repo_root.resolve():
142 worktree_name = "root"
143 else:
144 # Use directory name from workstack's work directory
145 worktree_name = wt.path.name
147 branch_to_worktree[wt.branch] = worktree_name
148 worktree_to_path[worktree_name] = wt.path
150 # Check if current path is within this worktree (handles subdirectories)
151 try:
152 current_path.relative_to(wt.path.resolve())
153 current_worktree = worktree_name
154 except ValueError:
155 # Not within this worktree
156 pass
158 return WorktreeMapping(
159 branch_to_worktree=branch_to_worktree,
160 worktree_to_path=worktree_to_path,
161 current_worktree=current_worktree,
162 )
165def _load_graphite_branch_graph(
166 ctx: WorkstackContext,
167 repo_root: Path,
168) -> BranchGraph | None:
169 """Load branch graph from Graphite cache.
171 Reads .git/.graphite_cache_persist JSON file and extracts parent-child
172 relationships between branches.
174 Args:
175 ctx: Workstack context with git operations
176 repo_root: Path to repository root
178 Returns:
179 BranchGraph if cache exists and is valid, None otherwise
180 """
181 # Get git common directory (handles both main repos and worktrees)
182 git_dir = ctx.git_ops.get_git_common_dir(repo_root)
183 if git_dir is None:
184 return None
186 # Check if Graphite cache file exists
187 cache_file = git_dir / ".graphite_cache_persist"
188 if not cache_file.exists():
189 return None
191 # Parse JSON
192 cache_data = json.loads(cache_file.read_text(encoding="utf-8"))
193 branches_data = cache_data.get("branches", [])
195 # Build relationship maps
196 parent_of: dict[str, str] = {}
197 children_of: dict[str, list[str]] = {}
198 trunk_branches: list[str] = []
200 for branch_name, info in branches_data:
201 parent = info.get("parentBranchName")
202 children = info.get("children", [])
203 is_trunk = info.get("validationResult") == "TRUNK"
205 # Record parent relationship
206 if parent:
207 parent_of[branch_name] = parent
209 # Record children
210 children_of[branch_name] = children
212 # Record trunk branches
213 if is_trunk or parent is None:
214 trunk_branches.append(branch_name)
216 return BranchGraph(
217 parent_of=parent_of,
218 children_of=children_of,
219 trunk_branches=trunk_branches,
220 )
223def _filter_graph_to_active_branches(
224 graph: BranchGraph,
225 active_branches: set[str],
226) -> BranchGraph:
227 """Filter branch graph to ONLY include branches with active worktrees.
229 This removes branches without worktrees from the graph while preserving
230 the tree structure. Only active branches and their relationships are kept.
232 Args:
233 graph: Full branch graph from Graphite cache
234 active_branches: Set of branch names that have worktrees
236 Returns:
237 Filtered BranchGraph containing only active branches
239 Example:
240 Input graph: main -> [feature-a, feature-b -> feature-b-2]
241 Active branches: {main, feature-a}
242 Output graph: main -> [feature-a]
243 (feature-b and feature-b-2 are removed)
244 """
245 filtered_parent_of: dict[str, str] = {}
246 filtered_children_of: dict[str, list[str]] = {}
247 filtered_trunk: list[str] = []
249 for branch in active_branches:
250 # Keep parent relationship only if branch is active
251 if branch in graph.parent_of:
252 filtered_parent_of[branch] = graph.parent_of[branch]
254 # Keep only children that are also active
255 if branch in graph.children_of:
256 active_children = [
257 child for child in graph.children_of[branch] if child in active_branches
258 ]
259 if active_children:
260 filtered_children_of[branch] = active_children
262 # Keep trunk status if active
263 if branch in graph.trunk_branches:
264 filtered_trunk.append(branch)
266 return BranchGraph(
267 parent_of=filtered_parent_of,
268 children_of=filtered_children_of,
269 trunk_branches=filtered_trunk,
270 )
273def _build_tree_from_graph(
274 graph: BranchGraph,
275 mapping: WorktreeMapping,
276) -> list[TreeNode]:
277 """Build TreeNode structure from filtered branch graph.
279 Recursively builds tree nodes starting from trunk branches, following
280 parent-child relationships to create the full tree structure.
282 Args:
283 graph: Filtered graph containing only active branches
284 mapping: Worktree mapping for annotations
286 Returns:
287 List of root TreeNode objects (one per trunk branch)
288 """
290 def build_node(branch: str) -> TreeNode:
291 """Recursively build a tree node and its children."""
292 worktree_name = mapping.branch_to_worktree[branch]
293 is_current = worktree_name == mapping.current_worktree
295 # Recursively build children
296 children_branches = graph.children_of.get(branch, [])
297 children = [build_node(child) for child in children_branches]
299 return TreeNode(
300 branch_name=branch,
301 worktree_name=worktree_name,
302 children=children,
303 is_current=is_current,
304 )
306 # Build tree starting from trunk branches
307 return [build_node(trunk) for trunk in graph.trunk_branches]
310def render_tree(roots: list[TreeNode]) -> str:
311 """Render tree structure as ASCII art with Unicode box-drawing characters.
313 Uses Unicode box-drawing characters:
314 - ├─ for middle children (branch continues below)
315 - └─ for last child (no more branches below)
316 - │ for continuation lines (shows vertical connection)
318 Args:
319 roots: List of root TreeNode objects
321 Returns:
322 Multi-line string with rendered tree
324 Example:
325 Input:
326 TreeNode("main", "root", [
327 TreeNode("feature-a", "feature-a", []),
328 TreeNode("feature-b", "feature-b", [])
329 ])
331 Output:
332 main [@root]
333 ├─ feature-a [@feature-a]
334 └─ feature-b [@feature-b]
335 """
336 lines: list[str] = []
338 def render_node(node: TreeNode, prefix: str, is_last: bool, is_root: bool) -> None:
339 """Recursively render a node and its children.
341 Args:
342 node: TreeNode to render
343 prefix: Prefix string for indentation (contains │ and spaces)
344 is_last: True if this is the last child of its parent
345 is_root: True if this is a top-level root node
346 """
347 # Format current line
348 connector = "└─" if is_last else "├─"
349 branch_text = _format_branch_name(node.branch_name, node.is_current)
350 worktree_text = _format_worktree_annotation(node.worktree_name)
352 if is_root:
353 # Root node: no connector
354 line = f"{branch_text} {worktree_text}"
355 else:
356 # All other nodes get connectors
357 line = f"{prefix}{connector} {branch_text} {worktree_text}"
359 lines.append(line)
361 # Render children
362 if node.children:
363 # Determine prefix for children
364 # Build prefix based on whether this node is the last child of its parent
365 if prefix:
366 # Non-root node: extend existing prefix
367 # Add vertical bar if more siblings below, space otherwise
368 child_prefix = prefix + (" " if is_last else "│ ")
369 else:
370 # Root node's children: start with appropriate spacing
371 # Use spaces if this is last root, vertical bar otherwise
372 child_prefix = " " if is_last else "│ "
374 for i, child in enumerate(node.children):
375 is_last_child = i == len(node.children) - 1
376 render_node(child, child_prefix, is_last_child, is_root=False)
378 # Render all roots
379 for i, root in enumerate(roots):
380 is_last_root = i == len(roots) - 1
381 render_node(root, "", is_last_root, is_root=True)
383 return "\n".join(lines)
386def _format_branch_name(branch: str, is_current: bool) -> str:
387 """Format branch name with color.
389 Args:
390 branch: Branch name to format
391 is_current: True if this is the current worktree
393 Returns:
394 Colored branch name (bright green if current, normal otherwise)
395 """
396 if is_current:
397 return click.style(branch, fg="bright_green", bold=True)
398 else:
399 return branch
402def _format_worktree_annotation(worktree_name: str) -> str:
403 """Format worktree annotation [@name].
405 Args:
406 worktree_name: Name of the worktree
408 Returns:
409 Dimmed annotation text
410 """
411 return click.style(f"[@{worktree_name}]", fg="bright_black")