feat: support rebasing branches checked out in worktrees

PyGitUp/git_wrapper.py

@@ -144,9 +144,11 @@ def stash():
 
             stashed[0] = True
 
+        stash.suppress_pop = False
+
         yield stash
 
-        if stashed[0]:
+        if stashed[0] and not stash.suppress_pop:
             print(colored('unstashing', 'magenta'))
             try:
                 self._run('stash', 'pop')

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.in_progress_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 has an in-progress operation
+                if branch.name in self.in_progress_branches:
+                    print(colored('operation in progress', 'yellow'))
+                    self.states.append('operation 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,105 @@ 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 = {}
+        in_progress_branches = set()
+        try:
+            output = self.git._run('worktree', 'list', '--porcelain')
+        except GitError:
+            return worktree_map, in_progress_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
+                    if self._worktree_has_in_progress_op(current_path):
+                        in_progress_branches.add(branch_name)
+            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
+                        in_progress_branches.add(branch_name)
+
+        return worktree_map, in_progress_branches
+
+    def _get_worktree_meta_dir(self, worktree_path):
+        """Return the git metadata directory for a 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)
+        return os.path.realpath(meta_dir)
+
+    def _worktree_has_in_progress_op(self, worktree_path):
+        """Return True if the worktree has a cherry-pick, merge, or bisect in progress."""
+        meta_dir = self._get_worktree_meta_dir(worktree_path)
+        if not meta_dir:
+            return False
+        for marker in ('CHERRY_PICK_HEAD', 'MERGE_HEAD', 'BISECT_LOG'):
+            if os.path.isfile(os.path.join(meta_dir, marker)):
+                return True
+        return False
+
+    def _get_rebase_branch(self, worktree_path):
+        """Return the branch name if a rebase is in progress in the worktree."""
+        meta_dir = self._get_worktree_meta_dir(worktree_path)
+        if not meta_dir:
+            return None
+        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:
+            with worktree_git.stasher() as stash:
+                stash()
+                try:
+                    worktree_git.rebase(target)
+                except RebaseError:
+                    stash.suppress_pop = True
+                    raise
+
     def fetch(self):
         """
         Fetch the recent refs from the remotes.
@@ -369,23 +482,23 @@ def push(self):
             error.message = "`git push` failed"
             raise error
 
-    def log(self, branch, remote):
-        """ Call a log-command, if set by git-up.fetch.all. """
-        log_hook = self.settings['rebase.log-hook']
-
-        if log_hook:
-            def _escape_positional(value):
-                # Neutralize command substitution/backticks in branch names
-                return value.replace('$', r'\$').replace('`', r'\`')
-
-            branch_safe = _escape_positional(branch.name)
-            remote_safe = _escape_positional(remote.name)
-
-            if ON_WINDOWS:  # pragma: no cover
-                # Running a string in CMD from Python is not that easy on
-                # Windows. Running 'cmd /C log_hook' produces problems when
-                # using multiple statements or things like 'echo'. Therefore,
-                # we write the string to a bat file and execute it.
+    def log(self, branch, remote):
+        """ Call a log-command, if set by git-up.fetch.all. """
+        log_hook = self.settings['rebase.log-hook']
+
+        if log_hook:
+            def _escape_positional(value):
+                # Neutralize command substitution/backticks in branch names
+                return value.replace('$', r'\$').replace('`', r'\`')
+
+            branch_safe = _escape_positional(branch.name)
+            remote_safe = _escape_positional(remote.name)
+
+            if ON_WINDOWS:  # pragma: no cover
+                # Running a string in CMD from Python is not that easy on
+                # Windows. Running 'cmd /C log_hook' produces problems when
+                # using multiple statements or things like 'echo'. Therefore,
+                # we write the string to a bat file and execute it.
 
                 # In addition, we replace occurrences of $1 with %1 and so forth
                 # in case the user is used to Bash or sh.
@@ -418,16 +531,16 @@ def _escape_positional(value):
 
                 # Clean up file
                 os.remove(bat_file.name)
-            else:  # pragma: no cover
-                # Run log_hook via 'shell -c'
-                # Disable globbing and word-splitting to keep $1/$2 safe
-                state = subprocess.call(
-                    ['sh', '-c', 'set -f; IFS=; ' + log_hook,
-                     'git-up', branch_safe, remote_safe]
-                )
-
-            if self.testing:
-                assert state == 0, 'log_hook returned != 0'
+            else:  # pragma: no cover
+                # Run log_hook via 'shell -c'
+                # Disable globbing and word-splitting to keep $1/$2 safe
+                state = subprocess.call(
+                    ['sh', '-c', 'set -f; IFS=; ' + log_hook,
+                     'git-up', branch_safe, remote_safe]
+                )
+
+            if self.testing:
+                assert state == 0, 'log_hook returned != 0'
 
     def version_info(self):
         """ Tell, what version we're running at and if it's up to date. """

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