Add --stash options to pull and rebase

docs/src/worktrees.md

@@ -309,6 +309,10 @@ grm wt pull --rebase
 [✔] my-cool-branch: Done
 ```
 
+As noted, this will fail if there are any local changes in your worktree. If you
+want to stash these changes automatically before the pull (and unstash them
+afterwards), use the `--stash` option.
+
 This will rebase your changes onto the upstream branch. This is mainly helpful
 for persistent branches that change on the remote side.
 
@@ -346,6 +350,10 @@ run two commands.
 I understand that the UX is not the most intuitive. If you can think of an
 improvement, please let me know (e.g. via an GitHub issue)!
 
+As with `pull`, `rebase` will also refuse to run when there are changes in your
+worktree. And you can also use the `--stash` option to stash/unstash changes
+automatically.
+
 ### Manual access
 
 GRM isn't doing any magic, it's just git under the hood. If you need to have access

e2e_tests/test_worktree_fetch.py

@@ -2,6 +2,8 @@
 
 from helpers import *
 
+import re
+
 import pytest
 import git
 
@@ -51,7 +53,9 @@ def test_worktree_fetch():
 
 @pytest.mark.parametrize("rebase", [True, False])
 @pytest.mark.parametrize("ffable", [True, False])
-def test_worktree_pull(rebase, ffable):
+@pytest.mark.parametrize("has_changes", [True, False])
+@pytest.mark.parametrize("stash", [True, False])
+def test_worktree_pull(rebase, ffable, has_changes, stash):
     with TempGitRepositoryWorktree() as (base_dir, root_commit):
         with TempGitFileRemote() as (remote_path, _remote_sha):
             shell(
@@ -94,51 +98,79 @@ def test_worktree_pull(rebase, ffable):
                     """
                     )
 
+                if has_changes:
+                    shell(
+                        f"""
+                        cd {base_dir}/master
+                        echo change >> root-commit-in-worktree-1
+                        echo uncommitedchange > uncommitedchange
+                    """
+                    )
+
                 args = ["wt", "pull"]
                 if rebase:
                     args += ["--rebase"]
+                if stash:
+                    args += ["--stash"]
                 cmd = grm(args, cwd=base_dir)
-                assert cmd.returncode == 0
+                if has_changes and not stash:
-
+                    assert cmd.returncode != 0
-                assert repo.commit("upstream/master").hexsha == remote_commit
+                    assert re.match(r".*master.*contains changes.*", cmd.stderr)
-                assert repo.commit("origin/master").hexsha == root_commit
-                assert (
-                    repo.commit("master").hexsha != repo.commit("origin/master").hexsha
-                )
-
-                if not rebase:
-                    if ffable:
-                        assert (
-                            repo.commit("master").hexsha
-                            != repo.commit("origin/master").hexsha
-                        )
-                        assert (
-                            repo.commit("master").hexsha
-                            == repo.commit("upstream/master").hexsha
-                        )
-                        assert repo.commit("upstream/master").hexsha == remote_commit
-                    else:
-                        assert "cannot be fast forwarded" in cmd.stderr
-                        assert (
-                            repo.commit("master").hexsha
-                            != repo.commit("origin/master").hexsha
-                        )
-                        assert repo.commit("master").hexsha != remote_commit
-                        assert repo.commit("upstream/master").hexsha == remote_commit
                 else:
-                    if ffable:
+                    assert repo.commit("upstream/master").hexsha == remote_commit
-                        assert (
+                    assert repo.commit("origin/master").hexsha == root_commit
-                            repo.commit("master").hexsha
+                    assert (
-                            != repo.commit("origin/master").hexsha
+                        repo.commit("master").hexsha
-                        )
+                        != repo.commit("origin/master").hexsha
-                        assert (
+                    )
-                            repo.commit("master").hexsha
+                    if has_changes:
-                            == repo.commit("upstream/master").hexsha
+                        assert ["uncommitedchange"] == repo.untracked_files
-                        )
+                        assert repo.is_dirty()
-                        assert repo.commit("upstream/master").hexsha == remote_commit
                     else:
-                        assert (
+                        assert not repo.is_dirty()
-                            repo.commit("master").message.strip()
+
-                            == "local-commit-in-master"
+                    if not rebase:
-                        )
+                        if ffable:
-                        assert repo.commit("master~1").hexsha == remote_commit
+                            assert cmd.returncode == 0
+                            assert (
+                                repo.commit("master").hexsha
+                                != repo.commit("origin/master").hexsha
+                            )
+                            assert (
+                                repo.commit("master").hexsha
+                                == repo.commit("upstream/master").hexsha
+                            )
+                            assert (
+                                repo.commit("upstream/master").hexsha == remote_commit
+                            )
+                        else:
+                            assert cmd.returncode != 0
+                            assert "cannot be fast forwarded" in cmd.stderr
+                            assert (
+                                repo.commit("master").hexsha
+                                != repo.commit("origin/master").hexsha
+                            )
+                            assert repo.commit("master").hexsha != remote_commit
+                            assert (
+                                repo.commit("upstream/master").hexsha == remote_commit
+                            )
+                    else:
+                        assert cmd.returncode == 0
+                        if ffable:
+                            assert (
+                                repo.commit("master").hexsha
+                                != repo.commit("origin/master").hexsha
+                            )
+                            assert (
+                                repo.commit("master").hexsha
+                                == repo.commit("upstream/master").hexsha
+                            )
+                            assert (
+                                repo.commit("upstream/master").hexsha == remote_commit
+                            )
+                        else:
+                            assert (
+                                repo.commit("master").message.strip()
+                                == "local-commit-in-master"
+                            )
+                            assert repo.commit("master~1").hexsha == remote_commit

e2e_tests/test_worktree_rebase.py

@@ -2,15 +2,18 @@
 
 from helpers import *
 
-import pytest
+import re
 
+import pytest
 import git
 
 
 @pytest.mark.parametrize("pull", [True, False])
 @pytest.mark.parametrize("rebase", [True, False])
 @pytest.mark.parametrize("ffable", [True, False])
-def test_worktree_rebase(pull, rebase, ffable):
+@pytest.mark.parametrize("has_changes", [True, False])
+@pytest.mark.parametrize("stash", [True, False])
+def test_worktree_rebase(pull, rebase, ffable, has_changes, stash):
     with TempGitRepositoryWorktree() as (base_dir, _root_commit):
         with open(os.path.join(base_dir, "grm.toml"), "w") as f:
             f.write('persistent_branches = ["mybasebranch"]')
@@ -83,6 +86,14 @@ def test_worktree_rebase(pull, rebase, ffable):
             """
             )
 
+        if has_changes:
+            shell(
+                f"""
+                cd {base_dir}/myfeatbranch
+                echo uncommitedchange > uncommitedchange
+            """
+            )
+
         grm(["wt", "delete", "--force", "tmp"], cwd=base_dir)
 
         repo = git.Repo(f"{base_dir}/.git-main-working-tree")
@@ -133,17 +144,23 @@ def test_worktree_rebase(pull, rebase, ffable):
             args += ["--pull"]
         if rebase:
             args += ["--rebase"]
+        if stash:
+            args += ["--stash"]
         cmd = grm(args, cwd=base_dir)
 
-        print(args)
         if rebase and not pull:
             assert cmd.returncode != 0
             assert len(cmd.stderr) != 0
+        elif has_changes and not stash:
+            assert cmd.returncode != 0
+            assert re.match(r".*myfeatbranch.*contains changes.*", cmd.stderr)
         else:
-            assert cmd.returncode == 0
             repo = git.Repo(f"{base_dir}/myfeatbranch")
+            if has_changes:
+                assert ["uncommitedchange"] == repo.untracked_files
             if pull:
                 if rebase:
+                    assert cmd.returncode == 0
                     if ffable:
                         assert (
                             repo.commit("HEAD").message.strip()
@@ -190,6 +207,7 @@ def test_worktree_rebase(pull, rebase, ffable):
                         assert repo.commit("HEAD~6").message.strip() == "commit-root"
                 else:
                     if ffable:
+                        assert cmd.returncode == 0
                         assert (
                             repo.commit("HEAD").message.strip()
                             == "commit-in-feat-remote"
@@ -208,6 +226,7 @@ def test_worktree_rebase(pull, rebase, ffable):
                         )
                         assert repo.commit("HEAD~4").message.strip() == "commit-root"
                     else:
+                        assert cmd.returncode != 0
                         assert (
                             repo.commit("HEAD").message.strip()
                             == "commit-in-feat-local-no-ff"
@@ -226,6 +245,7 @@ def test_worktree_rebase(pull, rebase, ffable):
                         )
                         assert repo.commit("HEAD~4").message.strip() == "commit-root"
             else:
+                assert cmd.returncode == 0
                 if ffable:
                     assert repo.commit("HEAD").message.strip() == "commit-in-feat-local"
                     assert (

src/grm/cmd.rs

@@ -132,6 +132,8 @@ pub struct WorktreeFetchArgs {}
 pub struct WorktreePullArgs {
     #[clap(long = "--rebase", help = "Perform a rebase instead of a fast-forward")]
     pub rebase: bool,
+    #[clap(long = "--stash", help = "Stash & unstash changes before & after pull")]
+    pub stash: bool,
 }
 
 #[derive(Parser)]
@@ -140,6 +142,11 @@ pub struct WorktreeRebaseArgs {
     pub pull: bool,
     #[clap(long = "--rebase", help = "Perform a rebase when doing a pull")]
     pub rebase: bool,
+    #[clap(
+        long = "--stash",
+        help = "Stash & unstash changes before & after rebase"
+    )]
+    pub stash: bool,
 }
 
 pub fn parse() -> Opts {

src/grm/main.rs

@@ -350,26 +350,27 @@ fn main() {
                         process::exit(1);
                     });
 
+                    let mut failures = false;
                     for worktree in repo.get_worktrees().unwrap_or_else(|error| {
                         print_error(&format!("Error getting worktrees: {}", error));
                         process::exit(1);
                     }) {
-                        if let Some(warning) =
+                        if let Some(warning) = worktree
-                            worktree
+                            .forward_branch(args.rebase, args.stash)
-                                .forward_branch(args.rebase)
+                            .unwrap_or_else(|error| {
-                                .unwrap_or_else(|error| {
+                                print_error(&format!("Error updating worktree branch: {}", error));
-                                    print_error(&format!(
+                                process::exit(1);
-                                        "Error updating worktree branch: {}",
+                            })
-                                        error
-                                    ));
-                                    process::exit(1);
-                                })
                         {
                             print_warning(&format!("{}: {}", worktree.name(), warning));
+                            failures = true;
                         } else {
                             print_success(&format!("{}: Done", worktree.name()));
                         }
                     }
+                    if failures {
+                        process::exit(1);
+                    }
                 }
                 cmd::WorktreeAction::Rebase(args) => {
                     if args.rebase && !args.pull {
@@ -406,10 +407,12 @@ fn main() {
                         process::exit(1);
                     });
 
+                    let mut failures = false;
+
                     for worktree in &worktrees {
                         if args.pull {
                             if let Some(warning) = worktree
-                                .forward_branch(args.rebase)
+                                .forward_branch(args.rebase, args.stash)
                                 .unwrap_or_else(|error| {
                                     print_error(&format!(
                                         "Error updating worktree branch: {}",
@@ -418,28 +421,29 @@ fn main() {
                                     process::exit(1);
                                 })
                             {
+                                failures = true;
                                 print_warning(&format!("{}: {}", worktree.name(), warning));
                             }
                         }
                     }
 
                     for worktree in &worktrees {
-                        if let Some(warning) =
+                        if let Some(warning) = worktree
-                            worktree
+                            .rebase_onto_default(&config, args.stash)
-                                .rebase_onto_default(&config)
+                            .unwrap_or_else(|error| {
-                                .unwrap_or_else(|error| {
+                                print_error(&format!("Error rebasing worktree branch: {}", error));
-                                    print_error(&format!(
+                                process::exit(1);
-                                        "Error rebasing worktree branch: {}",
+                            })
-                                        error
-                                    ));
-                                    process::exit(1);
-                                })
                         {
+                            failures = true;
                             print_warning(&format!("{}: {}", worktree.name(), warning));
                         } else {
                             print_success(&format!("{}: Done", worktree.name()));
                         }
                     }
+                    if failures {
+                        process::exit(1);
+                    }
                 }
             }
         }

src/repo.rs

@@ -181,17 +181,30 @@ impl Worktree {
         &self.name
     }
 
-    pub fn forward_branch(&self, rebase: bool) -> Result<Option<String>, String> {
+    pub fn forward_branch(&self, rebase: bool, stash: bool) -> Result<Option<String>, String> {
         let repo = Repo::open(Path::new(&self.name), false)
             .map_err(|error| format!("Error opening worktree: {}", error))?;
 
         if let Ok(remote_branch) = repo.find_local_branch(&self.name)?.upstream() {
             let status = repo.status(false)?;
+            let mut stashed_changes = false;
 
             if !status.clean() {
-                return Ok(Some(String::from("Worktree contains changes")));
+                if stash {
+                    repo.stash()?;
+                    stashed_changes = true;
+                } else {
+                    return Ok(Some(String::from("Worktree contains changes")));
+                }
             }
 
+            let unstash = || -> Result<(), String> {
+                if stashed_changes {
+                    repo.stash_pop()?;
+                }
+                Ok(())
+            };
+
             let remote_annotated_commit = repo
                 .0
                 .find_annotated_commit(remote_branch.commit()?.id().0)
@@ -231,6 +244,7 @@ impl Worktree {
                             continue;
                         }
                         rebase.abort().map_err(convert_libgit2_error)?;
+                        unstash()?;
                         return Err(convert_libgit2_error(error));
                     }
                 }
@@ -243,9 +257,11 @@ impl Worktree {
                     .map_err(convert_libgit2_error)?;
 
                 if analysis.is_up_to_date() {
+                    unstash()?;
                     return Ok(None);
                 }
                 if !analysis.is_fast_forward() {
+                    unstash()?;
                     return Ok(Some(String::from("Worktree cannot be fast forwarded")));
                 }
 
@@ -257,15 +273,18 @@ impl Worktree {
                     )
                     .map_err(convert_libgit2_error)?;
             }
+            unstash()?;
         } else {
             return Ok(Some(String::from("No remote branch to rebase onto")));
         };
+
         Ok(None)
     }
 
     pub fn rebase_onto_default(
         &self,
         config: &Option<WorktreeRootConfig>,
+        stash: bool,
     ) -> Result<Option<String>, String> {
         let repo = Repo::open(Path::new(&self.name), false)
             .map_err(|error| format!("Error opening worktree: {}", error))?;
@@ -291,6 +310,25 @@ impl Worktree {
             },
         };
 
+        let status = repo.status(false)?;
+        let mut stashed_changes = false;
+
+        if !status.clean() {
+            if stash {
+                repo.stash()?;
+                stashed_changes = true;
+            } else {
+                return Ok(Some(String::from("Worktree contains changes")));
+            }
+        }
+
+        let unstash = || -> Result<(), String> {
+            if stashed_changes {
+                repo.stash_pop()?;
+            }
+            Ok(())
+        };
+
         let base_branch = repo.find_local_branch(&default_branch_name)?;
         let base_annotated_commit = repo
             .0
@@ -330,11 +368,13 @@ impl Worktree {
                     continue;
                 }
                 rebase.abort().map_err(convert_libgit2_error)?;
+                unstash()?;
                 return Err(convert_libgit2_error(error));
             }
         }
 
         rebase.finish(None).map_err(convert_libgit2_error)?;
+        unstash()?;
         Ok(None)
     }
 }
@@ -456,6 +496,35 @@ impl Repo {
         }
     }
 
+    pub fn stash(&self) -> Result<(), String> {
+        let head_branch = self.head_branch()?;
+        let head = head_branch.commit()?;
+        let author = head.author();
+
+        // This is honestly quite horrible. The problem is that all stash operations expect a
+        // mutable reference (as they, well, mutate the repo after all). But we are heavily using
+        // immutable references a lot with this struct. I'm really not sure how to best solve this.
+        // Right now, we just open the repo AGAIN. It is safe, as we are only accessing the stash
+        // with the second reference, so there are no cross effects. But it just smells. Also,
+        // using `unwrap()` here as we are already sure that the repo is openable(?).
+        let mut repo = Repo::open(self.0.path(), false).unwrap();
+        repo.0
+            .stash_save2(&author, None, Some(git2::StashFlags::INCLUDE_UNTRACKED))
+            .map_err(convert_libgit2_error)?;
+        Ok(())
+    }
+
+    pub fn stash_pop(&self) -> Result<(), String> {
+        let mut repo = Repo::open(self.0.path(), false).unwrap();
+        repo.0
+            .stash_pop(
+                0,
+                Some(git2::StashApplyOptions::new().reinstantiate_index()),
+            )
+            .map_err(convert_libgit2_error)?;
+        Ok(())
+    }
+
     pub fn rename_remote(&self, remote: &RemoteHandle, new_name: &str) -> Result<(), String> {
         let failed_refspecs = self
             .0
@@ -1253,6 +1322,10 @@ impl Commit<'_> {
     pub fn id(&self) -> Oid {
         Oid(self.0.id())
     }
+
+    pub(self) fn author(&self) -> git2::Signature {
+        self.0.author()
+    }
 }
 
 impl<'a> Branch<'a> {