Coverage for src/workstack/cli/graphite.py: 92%
85 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-10-19 09:31 -0400
1"""Graphite integration for workstack.
3Graphite (https://graphite.dev) is a stacked git workflow tool that allows developers
4to manage dependent branches in linear stacks. This module reads graphite's internal
5metadata to display stack information for worktrees.
7For comprehensive gt mental model and command reference, see:
8.agent/GT_MENTAL_MODEL.md
10## What is Graphite?
12Graphite organizes branches into "stacks" - linear chains of dependent branches built
13on top of each other. For example:
15 main (trunk)
16 └─ feature/phase-1
17 └─ feature/phase-2
18 └─ feature/phase-3
20Each branch in the stack depends on its parent, making it easy to work on multiple
21related changes while keeping them in separate PRs.
23## Graphite Cache File
25Graphite maintains a persistent cache of branch relationships at:
26`.git/.graphite_cache_persist`
28This is a JSON file with the following structure:
30```json
31{
32 "branches": [
33 ["main", {
34 "validationResult": "TRUNK",
35 "children": ["feature/phase-1", "other-feature"]
36 }],
37 ["feature/phase-1", {
38 "parentBranchName": "main",
39 "children": ["feature/phase-2"]
40 }],
41 ["feature/phase-2", {
42 "parentBranchName": "feature/phase-1",
43 "children": []
44 }]
45 ]
46}
47```
49### Cache Structure Details
51- `branches`: Array of [branch_name, branch_metadata] tuples
52- `validationResult: "TRUNK"`: Marks the trunk branch (main/master)
53- `parentBranchName`: The parent branch in the stack (null for trunk)
54- `children`: Array of child branch names that branch off from this branch
56### Linear Stacks vs Branching
58While graphite supports branching (a trunk can have multiple children), individual
59stacks are typically linear. When displaying a stack for a specific branch, we show
60only the linear chain that branch belongs to:
62- Traverse DOWN: From current branch to trunk (collecting ancestors)
63- Traverse UP: From current branch, following the first child (collecting descendants)
65This gives us a linear chain even if the full graph has branches:
67```
68main
69 ├─ feature-a (not shown if we're on feature-b-1)
70 └─ feature-b-1
71 └─ feature-b-2 (linear stack: main → feature-b-1 → feature-b-2)
72```
74## Git Worktrees and Common Directory
76Git worktrees share the same `.git` directory structure. The actual `.git` directory
77may be in a different location than the worktree. We use `git rev-parse --git-common-dir`
78to find the shared git directory where `.graphite_cache_persist` is stored.
79"""
81import json
82from pathlib import Path
83from typing import Any, TypedDict
85from workstack.core.context import WorkstackContext
86from workstack.core.gitops import WorktreeInfo
89class BranchInfo(TypedDict):
90 """Metadata for a single branch in the graphite stack.
92 Fields:
93 parent: The parent branch name, or None if this is a trunk branch
94 children: List of child branch names that branch off from this branch
95 is_trunk: True if this is a trunk branch (main/master), False otherwise
96 """
98 parent: str | None
99 children: list[str]
100 is_trunk: bool
103def _load_graphite_cache(cache_file: Path) -> dict[str, Any]:
104 """Load and parse graphite cache file.
106 Args:
107 cache_file: Path to .graphite_cache_persist file
109 Returns:
110 Parsed cache data dictionary
112 Raises:
113 json.JSONDecodeError: If cache file is corrupted (fail-fast)
114 FileNotFoundError: If cache file doesn't exist
115 """
116 cache_data: dict[str, Any] = json.loads(cache_file.read_text(encoding="utf-8"))
117 return cache_data
120def _load_branch_info(ctx: WorkstackContext, repo_root: Path) -> dict[str, BranchInfo] | None:
121 """Load and parse branch info from Graphite cache.
123 Args:
124 ctx: Workstack context with git operations
125 repo_root: Path to the repository root (or worktree root)
127 Returns:
128 Dictionary mapping branch name to BranchInfo, or None if:
129 - Git command fails
130 - Graphite cache file doesn't exist
132 This helper consolidates the common cache loading logic used by
133 get_parent_branch, get_child_branches, and get_branch_stack.
134 """
135 git_dir = ctx.git_ops.get_git_common_dir(repo_root)
136 if git_dir is None:
137 return None
139 cache_file = git_dir / ".graphite_cache_persist"
140 if not cache_file.exists():
141 return None
143 cache_data = _load_graphite_cache(cache_file)
144 branches_data: list[Any] = cache_data.get("branches", [])
146 # Build parent-child relationship graph
147 branch_info: dict[str, BranchInfo] = {}
148 for branch_name, info in branches_data:
149 parent: str | None = info.get("parentBranchName")
150 children: list[str] = info.get("children", [])
151 is_trunk: bool = info.get("validationResult") == "TRUNK"
152 branch_info[branch_name] = BranchInfo(parent=parent, children=children, is_trunk=is_trunk)
154 return branch_info
157def get_branch_stack(ctx: WorkstackContext, repo_root: Path, branch: str) -> list[str] | None:
158 """Get the linear graphite stack for a given branch.
160 This function reads graphite's cache file and builds the linear chain of branches
161 that the given branch belongs to. The chain includes:
162 - All ancestor branches from current up to trunk
163 - All descendant branches from current down to the leaf
165 Args:
166 ctx: Workstack context with git operations
167 repo_root: Path to the repository root (or worktree root)
168 branch: Name of the branch to get the stack for
170 Returns:
171 List of branch names in the stack, ordered from trunk to leaf
172 (e.g., ["main", "feature-1", "feature-2", "feature-3"]).
173 Returns None if:
174 - Graphite cache file doesn't exist
175 - Git command fails
176 - Branch is not tracked by graphite
178 Algorithm:
179 1. Find the common git directory using ctx.git_ops.get_git_common_dir()
180 (This handles both main repos and worktrees correctly)
182 2. Load and parse `.graphite_cache_persist` JSON file
184 3. Build a parent-child relationship graph from the cache data
186 4. Traverse DOWN from current branch to trunk, collecting ancestors:
187 current → parent → grandparent → ... → trunk
189 5. Traverse UP from current branch, following the first child only:
190 current → child → grandchild → ... → leaf
192 6. Combine into linear chain: [trunk, ..., parent, current, child, ..., leaf]
194 Note on Linear vs Branching:
195 While graphite's cache can represent a tree structure (trunk with multiple
196 children), this function returns only the LINEAR chain that the given branch
197 belongs to. When a branch has multiple children, we follow only the first one.
199 This means if you have:
200 main
201 ├─ feature-a
202 └─ feature-b-1
203 └─ feature-b-2
205 And you call get_branch_stack(ctx, root, "feature-b-2"), you get:
206 ["main", "feature-b-1", "feature-b-2"]
208 Not: ["main", "feature-a", "feature-b-1", "feature-b-2"]
210 Future Enhancement: Could use smarter logic to follow the child that's actually
211 in the same stack as the target branch, rather than always taking the first child.
213 Example:
214 >>> stack = get_branch_stack(ctx, Path("/repo"), "feature/phase-2")
215 >>> print(stack)
216 ["main", "feature/phase-1", "feature/phase-2", "feature/phase-3"]
217 """
218 branch_info = _load_branch_info(ctx, repo_root)
219 if branch_info is None:
220 return None
222 # Check if the requested branch exists in graphite's cache
223 if branch not in branch_info:
224 return None
226 # Step 5: Traverse DOWN the stack to collect ancestors (current → parent → ... → trunk)
227 # In graphite terminology, "down" means towards the trunk/base
228 ancestors: list[str] = []
229 current = branch
230 while current in branch_info:
231 ancestors.append(current)
232 parent = branch_info[current]["parent"]
233 if parent is None or parent not in branch_info:
234 # Reached trunk (no parent) or parent not in cache
235 break
236 current = parent
238 # Reverse to get [trunk, ..., grandparent, parent, current]
239 ancestors.reverse()
241 # Step 6: Traverse UP the stack to collect descendants (current → child → ... → leaf)
242 # In graphite terminology, "up" means away from trunk towards the tip of the stack
243 # Only follow the first child to maintain a linear chain
244 descendants: list[str] = []
245 current = branch
246 while True:
247 children = branch_info[current]["children"]
248 if not children:
249 # Reached a leaf node (no children)
250 break
251 # For linear stacks, follow the first child
252 # NOTE: If there are multiple children (branching), we only follow the first one
253 # This keeps the stack linear, but may not show all related branches
254 first_child = children[0]
255 if first_child not in branch_info:
256 # Child exists in metadata but not in cache (shouldn't happen normally)
257 break
258 descendants.append(first_child)
259 current = first_child
261 # Step 7: Combine ancestors and descendants
262 # ancestors already includes the current branch, so we just append descendants
263 return ancestors + descendants
266def get_parent_branch(ctx: WorkstackContext, repo_root: Path, branch: str) -> str | None:
267 """Get the parent branch of a given branch in the Graphite stack.
269 Args:
270 ctx: Workstack context with git operations
271 repo_root: Path to the repository root (or worktree root)
272 branch: Name of the branch to get the parent for
274 Returns:
275 The parent branch name, or None if:
276 - Graphite cache file doesn't exist
277 - Git command fails
278 - Branch is not tracked by graphite
279 - Branch is at trunk (no parent)
281 Example:
282 >>> parent = get_parent_branch(ctx, Path("/repo"), "feature/phase-2")
283 >>> print(parent)
284 "feature/phase-1"
285 """
286 branch_info = _load_branch_info(ctx, repo_root)
287 if branch_info is None:
288 return None
290 # Check if the requested branch exists in graphite's cache
291 if branch not in branch_info:
292 return None
294 return branch_info[branch]["parent"]
297def get_child_branches(ctx: WorkstackContext, repo_root: Path, branch: str) -> list[str]:
298 """Get the child branches of a given branch in the Graphite stack.
300 Args:
301 ctx: Workstack context with git operations
302 repo_root: Path to the repository root (or worktree root)
303 branch: Name of the branch to get children for
305 Returns:
306 List of child branch names, or empty list if:
307 - Graphite cache file doesn't exist
308 - Git command fails
309 - Branch is not tracked by graphite
310 - Branch has no children (at tip of stack)
312 Example:
313 >>> children = get_child_branches(ctx, Path("/repo"), "feature/phase-1")
314 >>> print(children)
315 ["feature/phase-2", "feature/phase-2-alt"]
316 """
317 branch_info = _load_branch_info(ctx, repo_root)
318 if branch_info is None:
319 return []
321 # Check if the requested branch exists in graphite's cache
322 if branch not in branch_info:
323 return []
325 return branch_info[branch]["children"]
328def find_worktrees_containing_branch(
329 ctx: WorkstackContext,
330 repo_root: Path,
331 worktrees: list[WorktreeInfo],
332 target_branch: str,
333) -> list[WorktreeInfo]:
334 """Find all worktrees that have target_branch in their Graphite stack.
336 Args:
337 ctx: Workstack context with git operations
338 repo_root: Path to the repository root
339 worktrees: List of all worktrees from list_worktrees()
340 target_branch: Branch name to search for
342 Returns:
343 List of WorktreeInfo objects whose stacks contain target_branch.
344 Empty list if no worktrees contain the branch.
346 Algorithm:
347 1. For each worktree:
348 a. Get the worktree's checked-out branch
349 b. Skip worktrees with detached HEAD (branch=None)
350 c. Call get_branch_stack() to get the full stack for that branch
351 d. Check if target_branch is in that stack
352 e. If yes, add worktree to results
353 2. Return all matching worktrees
355 Example:
356 >>> worktrees = ctx.git_ops.list_worktrees(repo.root)
357 >>> matching = find_worktrees_containing_branch(ctx, repo.root, worktrees, "feature-2")
358 >>> print([wt.path for wt in matching])
359 [Path("/path/to/work/feature-work")]
360 """
361 matching_worktrees: list[WorktreeInfo] = []
363 for wt in worktrees:
364 # Skip worktrees with detached HEAD
365 if wt.branch is None:
366 continue
368 # Get the stack for this worktree's branch
369 stack = get_branch_stack(ctx, repo_root, wt.branch)
370 if stack is None:
371 # Graphite cache doesn't exist or branch not tracked
372 continue
374 # Check if target_branch is in this stack
375 if target_branch in stack:
376 matching_worktrees.append(wt)
378 return matching_worktrees
381def find_worktree_for_branch(worktrees: list[WorktreeInfo], branch: str) -> Path | None:
382 """Find the worktree path for a given branch.
384 Args:
385 worktrees: List of WorktreeInfo from ctx.git_ops.list_worktrees()
386 branch: Branch name to search for
388 Returns:
389 Path to the worktree if branch is checked out in a worktree,
390 None if no worktree exists for the branch.
392 Note:
393 This function expects WorktreeInfo objects with 'path' and 'branch' attributes.
395 Example:
396 >>> worktrees = ctx.git_ops.list_worktrees(repo.root)
397 >>> wt_path = find_worktree_for_branch(worktrees, "feature-1")
398 >>> print(wt_path)
399 Path("/path/to/work/feature-1")
400 """
401 for wt in worktrees:
402 if hasattr(wt, "branch") and wt.branch == branch:
403 return wt.path
404 return None