Keep local changes when checkout branch in worktree.

Signed-off-by: Linuxer Wang <linuxerwang@gmail.com>

3 files changed, 47 insertions, 0 deletions

diff --git a/options.go b/options.go
index 7c9e687..a3b14fe 100644
--- a/options.go
+++ b/options.go
@@ -242,6 +242,11 @@ type CheckoutOptions struct {
 	// Force, if true when switching branches, proceed even if the index or the
 	// working tree differs from HEAD. This is used to throw away local changes
 	Force bool
+	// Keep, if true when switching branches, local changes (the index or the
+	// working tree changes) will be kept so that they can be committed to the
+	// target branch. Force and Keep are mutually exclusive, should not be both
+	// set to true.
+	Keep bool
 }
 
 // Validate validates the fields and sets the default values.
diff --git a/worktree.go b/worktree.go
index dae40a3..1b10449 100644
--- a/worktree.go
+++ b/worktree.go
@@ -160,6 +160,8 @@ func (w *Worktree) Checkout(opts *CheckoutOptions) error {
 	ro := &ResetOptions{Commit: c, Mode: MergeReset}
 	if opts.Force {
 		ro.Mode = HardReset
+	} else if opts.Keep {
+		ro.Mode = SoftReset
 	}
 
 	if !opts.Hash.IsZero() && !opts.Create {
diff --git a/worktree_test.go b/worktree_test.go
index afedc91..045a76d 100644
--- a/worktree_test.go
+++ b/worktree_test.go
@@ -314,6 +314,46 @@ func (s *WorktreeSuite) TestCheckoutForce(c *C) {
 	c.Assert(entries, HasLen, 8)
 }
 
+func (s *WorktreeSuite) TestCheckoutKeep(c *C) {
+	w := &Worktree{
+		r:          s.Repository,
+		Filesystem: memfs.New(),
+	}
+
+	err := w.Checkout(&CheckoutOptions{
+		Force: true,
+	})
+	c.Assert(err, IsNil)
+
+	// Create a new branch and create a new file.
+	err = w.Checkout(&CheckoutOptions{
+		Branch: plumbing.NewBranchReferenceName("new-branch"),
+		Create: true,
+	})
+	c.Assert(err, IsNil)
+
+	w.Filesystem = memfs.New()
+	f, err := w.Filesystem.Create("new-file.txt")
+	c.Assert(err, IsNil)
+	_, err = f.Write([]byte("DUMMY"))
+	c.Assert(err, IsNil)
+	c.Assert(f.Close(), IsNil)
+
+	// Add the file to staging.
+	_, err = w.Add("new-file.txt")
+	c.Assert(err, IsNil)
+
+	// Switch branch to master, and verify that the new file was kept in staging.
+	err = w.Checkout(&CheckoutOptions{
+		Keep: true,
+	})
+	c.Assert(err, IsNil)
+
+	fi, err := w.Filesystem.Stat("new-file.txt")
+	c.Assert(err, IsNil)
+	c.Assert(fi.Size(), Equals, int64(5))
+}
+
 func (s *WorktreeSuite) TestCheckoutSymlink(c *C) {
 	if runtime.GOOS == "windows" {
 		c.Skip("git doesn't support symlinks by default in windows")