feat: support rebasing branches checked out in worktrees

PyGitUp/gitup.py

@@ -38,7 +38,7 @@
 
 # PyGitUp libs
 from PyGitUp.utils import execute, uniq, find
-from PyGitUp.git_wrapper import GitWrapper, GitError
+from PyGitUp.git_wrapper import GitWrapper, GitError, RebaseError
 
 ON_WINDOWS = sys.platform == 'win32'
 
@@ -185,6 +185,9 @@ def __init__(self, testing=False, sparse=False):
             self.git.status(porcelain=True, untracked_files='no').split('\n')
         )
 
+        # Build worktree map: branch name -> worktree path
+        self.worktree_map, self.mid_rebase_branches = self._build_worktree_map()
+
         # Load configuration
         self.settings = self.default_settings.copy()
         self.load_config()
@@ -247,6 +250,12 @@ def rebase_all_branches(self):
 
                     continue
 
+                # Skip branches whose worktree is mid-rebase
+                if branch.name in self.mid_rebase_branches:
+                    print(colored('rebase in progress', 'yellow'))
+                    self.states.append('rebase in progress')
+                    continue
+
                 # Get tracking branch
                 if target.is_local:
                     target = find(self.repo.branches,
@@ -291,7 +300,12 @@ def rebase_all_branches(self):
                     print()
 
                 self.log(branch, target)
-                if fast_fastforward:
+                worktree_path = self.worktree_map.get(branch.name)
+                if worktree_path:
+                    self._rebase_in_worktree(
+                        branch, target, worktree_path, fast_fastforward
+                    )
+                elif fast_fastforward:
                     branch.commit = target.commit
                 else:
                     stasher()
@@ -306,6 +320,106 @@ def rebase_all_branches(self):
                               'magenta'))
                 original_branch.checkout()
 
+    def _build_worktree_map(self):
+        """
+        Build a map of branch names to worktree paths.
+
+        This allows us to detect branches that are checked out in
+        separate worktrees, so we can rebase them in-place instead of
+        failing on checkout.
+        """
+        worktree_map = {}
+        mid_rebase_branches = set()
+        try:
+            output = self.git._run('worktree', 'list', '--porcelain')
+        except GitError:
+            return worktree_map, mid_rebase_branches
+
+        current_path = None
+        main_worktree = os.path.realpath(self.repo.working_dir)
+
+        for line in output.split('\n'):
+            line = line.rstrip('\r')
+            if line.startswith('worktree '):
+                current_path = line[len('worktree '):]
+            elif line.startswith('branch refs/heads/'):
+                branch_name = line[len('branch refs/heads/'):]
+                if current_path and \
+                        os.path.realpath(current_path) != main_worktree:
+                    worktree_map[branch_name] = current_path
+            elif line == 'detached' and current_path:
+                if os.path.realpath(current_path) != main_worktree:
+                    branch_name = self._get_rebase_branch(current_path)
+                    if branch_name:
+                        worktree_map[branch_name] = current_path
+                        mid_rebase_branches.add(branch_name)
+
+        return worktree_map, mid_rebase_branches
+
+    def _get_rebase_branch(self, worktree_path):
+        """Return the branch name if a rebase is in progress in the worktree."""
+        git_file = os.path.join(worktree_path, '.git')
+        if not os.path.isfile(git_file):
+            return None
+        with open(git_file, 'r') as f:
+            content = f.read().strip()
+        if not content.startswith('gitdir: '):
+            return None
+        meta_dir = content[len('gitdir: '):]
+        if not os.path.isabs(meta_dir):
+            meta_dir = os.path.join(worktree_path, meta_dir)
+        meta_dir = os.path.realpath(meta_dir)
+        for subdir in ('rebase-merge', 'rebase-apply'):
+            head_name_file = os.path.join(meta_dir, subdir, 'head-name')
+            if os.path.isfile(head_name_file):
+                with open(head_name_file, 'r') as f:
+                    ref = f.read().strip()
+                if ref.startswith('refs/heads/'):
+                    return ref[len('refs/heads/'):]
+        return None
+
+    def _rebase_in_worktree(self, branch, target, worktree_path,
+                            fast_forward):
+        """
+        Rebase or fast-forward a branch checked out in a worktree.
+
+        Instead of checking out the branch (which would fail), we operate
+        directly in the worktree directory where the branch is already
+        checked out.
+        """
+        worktree_repo = Repo(worktree_path, odbt=GitCmdObjectDB)
+        worktree_git = GitWrapper(worktree_repo)
+
+        if fast_forward:
+            worktree_git._run('merge', '--ff-only', target.name)
+        else:
+            # Stash worktree changes if needed
+            stashed = worktree_repo.is_dirty(submodules=False)
+            if stashed:
+                change_count = worktree_git.change_count
+                if change_count > 1:
+                    msg = f'stashing {change_count} changes in worktree'
+                else:
+                    msg = f'stashing {change_count} change in worktree'
+                print(colored(msg, 'magenta'))
+                worktree_git._run('stash')
+
+            try:
+                rebase_args = self.settings['rebase.arguments']
+                arguments = (
+                    ([rebase_args] if rebase_args else []) +
+                    [target.name]
+                )
+                try:
+                    worktree_git._run('rebase', *arguments)
+                except GitError as e:
+                    raise RebaseError(branch.name, target.name,
+                                      **e.__dict__)
+            finally:
+                if stashed:
+                    print(colored('unstashing in worktree', 'magenta'))
+                    worktree_git._run('stash', 'pop')
+
     def fetch(self):
         """
         Fetch the recent refs from the remotes.

PyGitUp/tests/test_worktree.py

@@ -0,0 +1,79 @@
+# System imports
+import os
+from os.path import join
+
+from git import *
+from PyGitUp.tests import basepath, init_master, update_file, write_file
+
+test_name = 'worktree-rebase'
+repo_path = join(basepath, test_name + os.sep)
+worktree_path = join(basepath, test_name + '-wt' + os.sep)
+
+
+def setup_module():
+    global master, repo
+
+    master_path, master = init_master(test_name)
+
+    # Prepare master repo
+    master.git.checkout(b=test_name)
+
+    # Clone to test repo
+    path = join(basepath, test_name)
+
+    master.clone(path, b=test_name)
+    repo = Repo(path, odbt=GitCmdObjectDB)
+
+    assert repo.working_dir == path
+
+    # Create a second branch that will be checked out in a worktree
+    repo.git.branch(test_name + '-wt', 'origin/' + test_name)
+
+    # Add the worktree with the second branch checked out
+    repo.git.worktree('add', worktree_path, test_name + '-wt')
+
+    # Set up tracking for the worktree branch
+    repo.git.branch('--set-upstream-to', 'origin/' + test_name,
+                     test_name + '-wt')
+
+    # Modify file in master to create something to rebase/fast-forward
+    update_file(master, test_name)
+
+
+def test_worktree():
+    """Run 'git up' with branches checked out in worktrees."""
+    os.chdir(repo_path)
+
+    # --- Fast-forward case ---
+    from PyGitUp.gitup import GitUp
+    gitup = GitUp(testing=True)
+    gitup.run()
+
+    assert 'fast-forwarding' in gitup.states
+
+    # The worktree branch should have been updated
+    assert (master.branches[test_name].commit ==
+            repo.branches[test_name + '-wt'].commit)
+
+    # --- Rebase case ---
+    # Make a local commit on the worktree branch so it diverges
+    wt_repo = Repo(worktree_path, odbt=GitCmdObjectDB)
+    wt_file = join(worktree_path, 'worktree_file.txt')
+    write_file(wt_file, 'worktree change')
+    wt_repo.index.add([wt_file])
+    wt_repo.index.commit('worktree commit')
+
+    # Make another commit on master so the branch diverges
+    update_file(master, test_name + ' second update')
+
+    gitup2 = GitUp(testing=True)
+    gitup2.run()
+
+    assert 'rebasing' in gitup2.states
+
+    # The worktree branch should contain the master commit
+    wt_repo = Repo(worktree_path, odbt=GitCmdObjectDB)
+    master_commit = master.branches[test_name].commit.hexsha
+    # Walk the worktree branch history to verify the master commit is there
+    wt_commits = [c.hexsha for c in wt_repo.iter_commits()]
+    assert master_commit in wt_commits