|
1 | 1 | import itertools |
2 | | -import os |
3 | 2 | from enum import auto |
4 | 3 | from typing import Dict, List, Optional, Type, Union |
5 | 4 |
|
@@ -55,48 +54,42 @@ def _update_worktrees_cache_after_checkout(self, checked_out_branch: LocalBranch |
55 | 54 | Only the current worktree's entry needs to be updated - linked worktrees |
56 | 55 | don't change when we checkout in a different worktree. |
57 | 56 | """ |
58 | | - current_worktree_path = self._git.get_root_dir() |
| 57 | + current_worktree_root_dir = self._git.get_current_worktree_root_dir() |
59 | 58 |
|
60 | | - # Remove the old branch entry for this worktree (if any) |
61 | | - # We need to find which branch was previously checked out in this worktree |
62 | | - for branch, path in self.__worktrees_cache.items(): |
63 | | - if path == current_worktree_path: # pragma: no branch |
64 | | - del self.__worktrees_cache[branch] |
| 59 | + for branch, path in self.__worktree_root_dir_for_branch.items(): |
| 60 | + if path == current_worktree_root_dir: # pragma: no branch |
| 61 | + del self.__worktree_root_dir_for_branch[branch] |
65 | 62 | break |
66 | 63 |
|
67 | | - # Add the new branch entry for this worktree |
68 | | - self.__worktrees_cache[checked_out_branch] = current_worktree_path |
| 64 | + self.__worktree_root_dir_for_branch[checked_out_branch] = current_worktree_root_dir |
69 | 65 |
|
70 | 66 | def _switch_to_branch_worktree( |
71 | 67 | self, |
72 | | - branch: LocalBranchShortName) -> None: |
| 68 | + target_branch: LocalBranchShortName) -> None: |
73 | 69 | """ |
74 | 70 | Switch to the worktree where the branch is checked out, or to main worktree if branch is not checked out. |
75 | 71 | This may involve changing the current working directory. |
76 | 72 | Updates the worktrees cache after checkout. |
77 | 73 | """ |
78 | | - worktree_path = self.__worktrees_cache.get(branch) |
79 | | - current_worktree_root = self._git.get_root_dir() |
| 74 | + target_worktree_root_dir = self.__worktree_root_dir_for_branch.get(target_branch) |
| 75 | + current_worktree_root_dir = self._git.get_current_worktree_root_dir() |
80 | 76 |
|
81 | | - if worktree_path is None: |
| 77 | + if target_worktree_root_dir is None: |
82 | 78 | # Branch is not checked out anywhere, need to checkout in main worktree |
83 | 79 | # Only cd if we're currently in a different worktree (linked worktree) |
84 | | - main_worktree_path = self._git.get_main_worktree_path() |
85 | | - if current_worktree_root != main_worktree_path: |
86 | | - print(f"Changing directory to main worktree at {bold(main_worktree_path)}") |
87 | | - os.chdir(main_worktree_path) |
88 | | - # checkout() below will flush all caches, including __root_dir |
89 | | - self._git.checkout(branch) |
| 80 | + main_worktree_root_dir = self._git.get_main_worktree_root_dir() |
| 81 | + if current_worktree_root_dir != main_worktree_root_dir: |
| 82 | + print(f"Changing directory to main worktree at {bold(main_worktree_root_dir)}") |
| 83 | + self._git.chdir(main_worktree_root_dir) |
| 84 | + self._git.checkout(target_branch) |
90 | 85 | # Update cache after checkout |
91 | | - self._update_worktrees_cache_after_checkout(branch) |
| 86 | + self._update_worktrees_cache_after_checkout(target_branch) |
92 | 87 | else: |
93 | 88 | # Branch is checked out in a worktree |
94 | 89 | # Only cd if we're in a different worktree |
95 | | - if current_worktree_root != worktree_path: |
96 | | - print(f"Changing directory to {bold(worktree_path)} worktree where {bold(branch)} is checked out") |
97 | | - os.chdir(worktree_path) |
98 | | - # Flush root dir cache after directory change so get_root_dir() returns the correct path |
99 | | - self._git.flush_root_dir_cache() |
| 90 | + if current_worktree_root_dir != target_worktree_root_dir: |
| 91 | + print(f"Changing directory to {bold(target_worktree_root_dir)} worktree where {bold(target_branch)} is checked out") |
| 92 | + self._git.chdir(target_worktree_root_dir) |
100 | 93 |
|
101 | 94 | def traverse( |
102 | 95 | self, |
@@ -140,10 +133,10 @@ def traverse( |
140 | 133 |
|
141 | 134 | # Store the initial directory for later restoration |
142 | 135 | initial_branch = nearest_remaining_branch = self._git.get_current_branch() |
143 | | - initial_worktree_root = self._git.get_root_dir() |
| 136 | + initial_worktree_root = self._git.get_current_worktree_root_dir() |
144 | 137 |
|
145 | 138 | # Fetch worktrees once at the start to avoid repeated git worktree list calls |
146 | | - self.__worktrees_cache: Dict[LocalBranchShortName, str] = self._git.get_worktrees() |
| 139 | + self.__worktree_root_dir_for_branch: Dict[LocalBranchShortName, str] = self._git.get_worktree_root_dirs_by_branch() |
147 | 140 |
|
148 | 141 | try: |
149 | 142 | if opt_start_from == TraverseStartFrom.ROOT: |
@@ -502,7 +495,7 @@ def traverse( |
502 | 495 | finally: |
503 | 496 | # Warn if the initial directory doesn't correspond to the final checked out branch's worktree |
504 | 497 | final_branch = self._git.get_current_branch() |
505 | | - final_worktree_path = self.__worktrees_cache.get(final_branch) |
| 498 | + final_worktree_path = self.__worktree_root_dir_for_branch.get(final_branch) |
506 | 499 | if final_worktree_path and initial_worktree_root != final_worktree_path: |
507 | 500 | # Final branch is checked out in a worktree different from where we started |
508 | 501 | warn( |
|
0 commit comments