Implement worktree support for traverse command

RELEASE_NOTES.md

@@ -2,6 +2,7 @@
 
 ## New in git-machete 3.37.1
 
+- improved: `traverse` now changes directory to the worktree where a branch is checked out, rather than failing (suggested by @lsierant)
 - fixed: use `simple` (rather than `none`) mode of squash merge detection in `git machete discover`, to keep parity with the IntelliJ plugin
 
 ## New in git-machete 3.37.0

docs/man/git-machete.1

@@ -2143,6 +2143,8 @@ Operations like \fBgit machete github anno\-prs\fP (\fBgit machete gitlab anno\-
 and \fBgit machete github checkout\-prs\fP (\fBgit machete gitlab checkout\-mrs\fP) add \fBrebase=no push=no\fP branch qualifiers
 when the current user is NOT the author of the PR/MR associated with that branch.
 .sp
+\fBNote on git worktrees:\fP if a branch is already checked out in another worktree, \fBtraverse\fP will change directory to that worktree rather than failing.
+.sp
 \fBOptions:\fP
 .INDENT 0.0
 .TP

docs/source/cli/traverse.rst

@@ -91,6 +91,8 @@ Operations like ``git machete github anno-prs`` (``git machete gitlab anno-mrs``
 and ``git machete github checkout-prs`` (``git machete gitlab checkout-mrs``) add ``rebase=no push=no`` branch qualifiers
 when the current user is NOT the author of the PR/MR associated with that branch.
 
+**Note on git worktrees:** if a branch is already checked out in another worktree, ``traverse`` will change directory to that worktree rather than failing.
+
 **Options:**
 
 -F, --fetch                    Fetch the remotes of all managed branches at the beginning of traversal (no ``git pull`` involved, only ``git fetch``).

git_machete/client/traverse.py

@@ -1,6 +1,6 @@
 import itertools
 from enum import auto
-from typing import List, Optional, Type, Union
+from typing import Dict, List, Optional, Type, Union
 
 from git_machete.annotation import Annotation, Qualifiers
 from git_machete.client.base import (ParsableEnum, PickRoot,
@@ -11,7 +11,7 @@
 from git_machete.git_operations import (GitContext, LocalBranchShortName,
                                         SyncToRemoteStatus)
 from git_machete.utils import (bold, flat_map, fmt, get_pretty_choices,
-                               get_right_arrow)
+                               get_right_arrow, warn)
 
 
 class TraverseReturnTo(ParsableEnum):
@@ -46,6 +46,51 @@ def from_string_or_branch(cls: Type['TraverseStartFrom'], value: str,
 
 
 class TraverseMacheteClient(MacheteClientWithCodeHosting):
+    def _update_worktrees_cache_after_checkout(self, checked_out_branch: LocalBranchShortName) -> None:
+        """
+        Update the worktrees cache after a checkout operation in the current worktree.
+        This avoids the need to re-fetch the full worktree list.
+
+        Only the current worktree's entry needs to be updated - linked worktrees
+        don't change when we checkout in a different worktree.
+        """
+        current_worktree_root_dir = self._git.get_current_worktree_root_dir()
+
+        for branch, path in self.__worktree_root_dir_for_branch.items():
+            if path == current_worktree_root_dir:  # pragma: no branch
+                del self.__worktree_root_dir_for_branch[branch]
+                break
+
+        self.__worktree_root_dir_for_branch[checked_out_branch] = current_worktree_root_dir
+
+    def _switch_to_branch_worktree(
+            self,
+            target_branch: LocalBranchShortName) -> None:
+        """
+        Switch to the worktree where the branch is checked out, or to main worktree if branch is not checked out.
+        This may involve changing the current working directory.
+        Updates the worktrees cache after checkout.
+        """
+        target_worktree_root_dir = self.__worktree_root_dir_for_branch.get(target_branch)
+        current_worktree_root_dir = self._git.get_current_worktree_root_dir()
+
+        if target_worktree_root_dir is None:
+            # Branch is not checked out anywhere, need to checkout in main worktree
+            # Only cd if we're currently in a different worktree (linked worktree)
+            main_worktree_root_dir = self._git.get_main_worktree_root_dir()
+            if current_worktree_root_dir != main_worktree_root_dir:
+                print(f"Changing directory to main worktree at {bold(main_worktree_root_dir)}")
+                self._git.chdir(main_worktree_root_dir)
+            self._git.checkout(target_branch)
+            # Update cache after checkout
+            self._update_worktrees_cache_after_checkout(target_branch)
+        else:
+            # Branch is checked out in a worktree
+            # Only cd if we're in a different worktree
+            if current_worktree_root_dir != target_worktree_root_dir:
+                print(f"Changing directory to {bold(target_worktree_root_dir)} worktree where {bold(target_branch)} is checked out")
+                self._git.chdir(target_worktree_root_dir)
+
     def traverse(
             self,
             *,
@@ -86,21 +131,26 @@ def traverse(
             self._init_code_hosting_client()
             current_user = self.code_hosting_client.get_current_user_login()
 
+        # Store the initial directory for later restoration
         initial_branch = nearest_remaining_branch = self._git.get_current_branch()
+        initial_worktree_root = self._git.get_current_worktree_root_dir()
+
+        # Fetch worktrees once at the start to avoid repeated git worktree list calls
+        self.__worktree_root_dir_for_branch: Dict[LocalBranchShortName, str] = self._git.get_worktree_root_dirs_by_branch()
 
         try:
             if opt_start_from == TraverseStartFrom.ROOT:
                 dest = self.root_branch_for(self._git.get_current_branch(), if_unmanaged=PickRoot.FIRST)
                 self._print_new_line(False)
                 print(f"Checking out the root branch ({bold(dest)})")
-                self._git.checkout(dest)
+                self._switch_to_branch_worktree(dest)
                 current_branch = dest
             elif opt_start_from == TraverseStartFrom.FIRST_ROOT:
                 # Note that we already ensured that there is at least one managed branch.
                 dest = self.managed_branches[0]
                 self._print_new_line(False)
                 print(f"Checking out the first root branch ({bold(dest)})")
-                self._git.checkout(dest)
+                self._switch_to_branch_worktree(dest)
                 current_branch = dest
             elif opt_start_from == TraverseStartFrom.HERE:
                 current_branch = self._git.get_current_branch()
@@ -110,7 +160,7 @@ def traverse(
                 self.expect_in_managed_branches(dest)
                 self._print_new_line(False)
                 print(f"Checking out branch {bold(dest)}")
-                self._git.checkout(dest)
+                self._switch_to_branch_worktree(dest)
                 current_branch = dest
             else:
                 raise UnexpectedMacheteException(f"Unexpected value for opt_start_from: {opt_start_from}")
@@ -187,7 +237,7 @@ def traverse(
                 if branch != current_branch and needs_any_action:
                     self._print_new_line(False)
                     print(f"Checking out {bold(branch)}")
-                    self._git.checkout(branch)
+                    self._switch_to_branch_worktree(branch)
                     current_branch = branch
                     self._print_new_line(False)
                     self.status(
@@ -411,9 +461,12 @@ def traverse(
                     break
 
             if opt_return_to == TraverseReturnTo.HERE:
-                self._git.checkout(initial_branch)
+                # Return to initial branch
+                # No point switching back to initial directory as cwd won't propagate back to the calling shell anyway
+                self._switch_to_branch_worktree(initial_branch)
             elif opt_return_to == TraverseReturnTo.NEAREST_REMAINING:
-                self._git.checkout(nearest_remaining_branch)
+                self._switch_to_branch_worktree(nearest_remaining_branch)
+                # For NEAREST_REMAINING, we stay in the worktree where the branch is checked out
             # otherwise opt_return_to == TraverseReturnTo.STAY, so no action is needed
 
             self._print_new_line(False)
@@ -440,4 +493,12 @@ def traverse(
                     f"The initial branch {bold(initial_branch)} has been slid out. "
                     f"Returned to nearest remaining managed branch {bold(nearest_remaining_branch)}")
         finally:
-            pass
+            # Warn if the initial directory doesn't correspond to the final checked out branch's worktree
+            final_branch = self._git.get_current_branch()
+            final_worktree_path = self.__worktree_root_dir_for_branch.get(final_branch)
+            if final_worktree_path and initial_worktree_root != final_worktree_path:
+                # Final branch is checked out in a worktree different from where we started
+                warn(
+                    f"branch {bold(final_branch)} is checked out in worktree at {bold(final_worktree_path)}\n"
+                    f"You may want to change directory with:\n"
+                    f"  `cd {final_worktree_path}`")

git_machete/generated_docs.py

@@ -1485,6 +1485,8 @@
         and `git machete github checkout-prs` (`git machete gitlab checkout-mrs`) add `rebase=no push=no` branch qualifiers
         when the current user is NOT the author of the PR/MR associated with that branch.
 
+        <b>Note on git worktrees:</b> if a branch is already checked out in another worktree, `traverse` will change directory to that worktree rather than failing.
+
         <b>Options:</b>
 
            <b>-F</b>, <b>--fetch</b>

git_machete/git_operations.py

@@ -219,6 +219,7 @@ def __init__(self) -> None:
         self.owner: Optional[Any] = None
 
         self.__git_version: Optional[Tuple[int, int, int]] = None
+        self.__main_worktree_root_dir: Optional[str] = None
         self.__main_worktree_git_dir: Optional[str] = None
         self.__current_worktree_root_dir: Optional[str] = None
         self.__current_worktree_git_dir: Optional[str] = None
@@ -256,8 +257,13 @@ def flush_caches(self) -> None:
         self.__short_commit_hash_by_revision_cached = {}
         self.__flush_current_worktree_caches()
 
+    def chdir(self, path: str) -> None:
+        os.chdir(path)
+        self.__flush_current_worktree_caches()
+
     def __flush_current_worktree_caches(self) -> None:
         self.__current_worktree_root_dir = None
+        self.__main_worktree_root_dir = None
 
     def _run_git(self, git_cmd: str, *args: str, flush_caches: bool, allow_non_zero: bool = False) -> int:
         exit_code = utils.run_cmd(*GIT_EXEC, git_cmd, *args)
@@ -573,6 +579,75 @@ def is_merge_in_progress(self) -> bool:
     def is_revert_in_progress(self) -> bool:
         return os.path.isfile(self.get_current_worktree_git_subpath("REVERT_HEAD"))
 
+    def get_worktree_root_dirs_by_branch(self) -> Dict[LocalBranchShortName, str]:
+        """
+        Returns a dict mapping branch names to their worktree paths.
+        The main worktree's current branch is included if it's on a branch.
+        Branches not checked out in any worktree are not in the dict.
+        Worktrees in detached HEAD state are not included (since they can't be looked up by branch name).
+        """
+        if self.get_git_version() < (2, 5):
+            # git worktree command was introduced in git 2.5
+            return {}
+
+        result = self._popen_git("worktree", "list", "--porcelain")
+        worktrees: Dict[LocalBranchShortName, str] = {}
+
+        current_worktree_path: Optional[str] = None
+        current_branch: Optional[LocalBranchShortName] = None
+        is_detached: bool = False
+
+        for line in result.stdout.strip().splitlines():
+            if line.startswith('worktree '):
+                current_worktree_path = line[len('worktree '):]
+            elif line.startswith('branch '):
+                # Format: "branch refs/heads/<branch-name>"
+                branch_ref = line[len('branch '):]
+                current_branch = LocalBranchFullName.of(branch_ref).to_short_name()
+                is_detached = False
+            elif line.startswith('detached'):
+                is_detached = True
+                current_branch = None
+            elif line == '':
+                # Empty line marks end of worktree entry
+                # Only add worktrees that have a branch (not in detached HEAD state)
+                if current_worktree_path and current_branch is not None and not is_detached:
+                    worktrees[current_branch] = current_worktree_path
+                current_worktree_path = None
+                current_branch = None
+                is_detached = False
+
+        # Handle last entry if no trailing newline
+        if current_worktree_path and current_branch is not None and not is_detached:
+            worktrees[current_branch] = current_worktree_path
+
+        return worktrees
+
+    def get_main_worktree_root_dir(self) -> str:
+        """
+        Returns the path to the main worktree (not a linked worktree).
+        """
+        if not self.__main_worktree_root_dir:
+            if self.get_git_version() < (2, 5):
+                # git worktree command was introduced in git 2.5
+                # If worktrees aren't supported, just return the current root dir
+                self.__main_worktree_root_dir = self.get_current_worktree_root_dir()
+            else:
+                # We can't rely on `git rev-parse --git-common-dir`
+                # since in some earlier supported versions of git
+                # this path is apparently printed in a faulty way when dealing with worktrees.
+                # Let's parse the output of of `git worktree list` instead.
+                result = self._popen_git("worktree", "list", "--porcelain")
+
+                # The first worktree in the list is always the main worktree
+                first_entry = result.stdout.splitlines()[0]
+                if first_entry.startswith('worktree '):
+                    self.__main_worktree_root_dir = first_entry[len('worktree '):]
+                else:
+                    raise UnexpectedMacheteException("Could not obtain the main worktree path from "
+                                                     f"`git worktree list --porcelain` output (first line: `{first_entry}`)")
+        return self.__main_worktree_root_dir
+
     def checkout(self, branch: LocalBranchShortName) -> None:
         self._run_git("checkout", "--quiet", branch, "--", flush_caches=True)