Implement worktree support for traverse command

- Add get_worktrees() and get_worktree_path() methods to GitOperations
- Modify traverse to cd into worktrees where branches are checked out
- Cache worktree information to avoid repeated git worktree list calls
- Add informative messages when changing directories to worktrees
- Add comprehensive tests for worktree handling including detached HEAD cases
- Add integration test for traverse with multiple worktrees

Fixes #1522

Author: PawelLipski

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

ci/tox/build-context/run.sh

@@ -11,7 +11,8 @@ env | sort | head -4
 
 set -x
 
-tox -e mypy,coverage
+tox -e mypy
+tox -e coverage -- -vv
 
 $PYTHON -m build
 

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/cli.py

@@ -503,12 +503,12 @@ def set_utils_global_variables(parsed_args: argparse.Namespace) -> None:
 
 def launch(orig_args: List[str]) -> None:
     try:
-        launch_interruptible(orig_args)
+        launch_internal(orig_args)
     except InteractionStopped:
         pass
 
 
-def launch_interruptible(orig_args: List[str]) -> None:
+def launch_internal(orig_args: List[str]) -> None:
     initial_current_directory: Optional[str] = utils.get_current_directory_or_none()
 
     try:

git_machete/client/traverse.py

@@ -1,6 +1,7 @@
 import itertools
+import os
 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 +12,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 +47,57 @@ 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_path = self._git.get_root_dir()
+
+        # Remove the old branch entry for this worktree (if any)
+        # We need to find which branch was previously checked out in this worktree
+        for branch, path in self.__worktrees_cache.items():
+            if path == current_worktree_path:  # pragma: no branch
+                del self.__worktrees_cache[branch]
+                break
+
+        # Add the new branch entry for this worktree
+        self.__worktrees_cache[checked_out_branch] = current_worktree_path
+
+    def _switch_to_branch_worktree(
+            self,
+            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.
+        """
+        worktree_path = self.__worktrees_cache.get(branch)
+        current_worktree_root = self._git.get_root_dir()
+
+        if worktree_path 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_path = self._git.get_main_worktree_path()
+            if current_worktree_root != main_worktree_path:
+                print(f"Changing directory to main worktree at {bold(main_worktree_path)}")
+                os.chdir(main_worktree_path)
+                # checkout() below will flush all caches, including __root_dir
+            self._git.checkout(branch)
+            # Update cache after checkout
+            self._update_worktrees_cache_after_checkout(branch)
+        else:
+            # Branch is checked out in a worktree
+            # Only cd if we're in a different worktree
+            if current_worktree_root != worktree_path:
+                print(f"Changing directory to {bold(worktree_path)} worktree where {bold(branch)} is checked out")
+                os.chdir(worktree_path)
+                # Flush root dir cache after directory change so get_root_dir() returns the correct path
+                self._git.flush_root_dir_cache()
+
     def traverse(
             self,
             *,
@@ -86,21 +138,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_root_dir()
+
+        # Fetch worktrees once at the start to avoid repeated git worktree list calls
+        self.__worktrees_cache: Dict[LocalBranchShortName, str] = self._git.get_worktrees()
 
         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 +167,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 +244,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 +468,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 +500,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.__worktrees_cache.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"Note: 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

@@ -220,6 +220,7 @@ def __init__(self) -> None:
 
         self.__git_version: Optional[Tuple[int, int, int]] = None
         self.__root_dir: Optional[str] = None
+        self.__main_worktree: Optional[str] = None
         self.__main_git_dir: Optional[str] = None
         self.__worktree_git_dir: Optional[str] = None
 
@@ -253,8 +254,14 @@ def flush_caches(self) -> None:
         self.__remote_branches_cached = None
         self.__remotes_cached = None
         self.__removed_from_remote = None
+        self.__root_dir = None
+        self.__main_worktree = None
         self.__short_commit_hash_by_revision_cached = {}
 
+    def flush_root_dir_cache(self) -> None:
+        """Flush only the cached root directory. Useful after changing directories with os.chdir()."""
+        self.__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)
         if flush_caches:
@@ -569,6 +576,83 @@ def is_merge_in_progress(self) -> bool:
     def is_revert_in_progress(self) -> bool:
         return os.path.isfile(self.get_worktree_git_subpath("REVERT_HEAD"))
 
+    def get_worktrees(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_worktree_path(self, branch: LocalBranchShortName) -> Optional[str]:
+        """
+        Returns the path to the worktree where the given branch is currently checked out,
+        or None if the branch is not checked out in any worktree.
+        """
+        worktrees = self.get_worktrees()
+        return worktrees.get(branch)
+
+    def get_main_worktree_path(self) -> str:
+        """
+        Returns the path to the main worktree (not a linked worktree).
+        """
+        if not self.__main_worktree:
+            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 = self.get_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 = 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
+
     def checkout(self, branch: LocalBranchShortName) -> None:
         self._run_git("checkout", "--quiet", branch, "--", flush_caches=True)
 

graphics/setup-sandbox-with-worktrees

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+
+set -e -o pipefail -u
+
+root_dir=$(realpath "${1-$HOME}")
+
+# First, set up the regular sandbox
+script_dir=$(cd "$(dirname "$0")" && pwd)
+"$script_dir/setup-sandbox" "$root_dir"
+
+# Now create worktrees
+cd "$root_dir/machete-sandbox"
+
+echo
+echo "Creating worktrees..."
+echo
+
+# Create worktrees for some branches
+git worktree add "$root_dir/machete-sandbox-worktrees/allow-ownership-link" allow-ownership-link
+echo "Created worktree for allow-ownership-link at $root_dir/machete-sandbox-worktrees/allow-ownership-link"
+
+git worktree add "$root_dir/machete-sandbox-worktrees/build-chain" build-chain
+echo "Created worktree for build-chain at $root_dir/machete-sandbox-worktrees/build-chain"
+
+git worktree add "$root_dir/machete-sandbox-worktrees/master" master
+echo "Created worktree for master at $root_dir/machete-sandbox-worktrees/master"
+
+# Create a worktree with detached HEAD at the tip of call-ws
+call_ws_sha=$(git rev-parse call-ws)
+git worktree add --detach "$root_dir/machete-sandbox-worktrees/detached-call-ws" "$call_ws_sha"
+echo "Created worktree with detached HEAD at $root_dir/machete-sandbox-worktrees/detached-call-ws (commit $call_ws_sha from call-ws)"
+
+# Create another worktree with detached HEAD at the tip of hotfix/add-trigger
+hotfix_sha=$(git rev-parse hotfix/add-trigger)
+git worktree add --detach "$root_dir/machete-sandbox-worktrees/detached-hotfix" "$hotfix_sha"
+echo "Created worktree with detached HEAD at $root_dir/machete-sandbox-worktrees/detached-hotfix (commit $hotfix_sha from hotfix/add-trigger)"
+
+echo
+echo "Worktree list:"
+git worktree list
+echo
+echo
+git machete status
+echo
+echo

tests/mockers_git_repository.py

@@ -153,3 +153,13 @@ def set_git_config_key(key: str, value: str) -> None:
 
 def unset_git_config_key(key: str) -> None:
     execute(f'git config --unset {key}')
+
+
+def add_worktree(branch: str) -> str:
+    """
+    Create a new worktree for the specified branch in a temporary directory.
+    Returns the path to the created worktree.
+    """
+    worktree_path = mkdtemp()
+    execute(f"git worktree add -f {worktree_path} {branch}")
+    return worktree_path