Merge pull request #493 from openmetagame/git-restore-for-pr-2022

Add RestoreStaged to Worktree that mimics the behaviour of git restore --staged <file>...

5 files changed, 382 insertions, 9 deletions

diff --git a/_examples/common_test.go b/_examples/common_test.go
index 5e3f753..7b6dbfb 100644
--- a/_examples/common_test.go
+++ b/_examples/common_test.go
@@ -30,6 +30,7 @@ var args = map[string][]string{
 	"progress":                   {defaultURL, tempFolder()},
 	"pull":                       {createRepositoryWithRemote(tempFolder(), defaultURL)},
 	"push":                       {setEmptyRemote(cloneRepository(defaultURL, tempFolder()))},
+	"restore":                    {cloneRepository(defaultURL, tempFolder())},
 	"revision":                   {cloneRepository(defaultURL, tempFolder()), "master~2^"},
 	"sha256":                     {tempFolder()},
 	"showcase":                   {defaultURL, tempFolder()},
diff --git a/_examples/restore/main.go b/_examples/restore/main.go
new file mode 100644
index 0000000..8016b06
--- /dev/null
+++ b/_examples/restore/main.go
@@ -0,0 +1,103 @@
+package main
+
+import (
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/go-git/go-git/v5"
+	. "github.com/go-git/go-git/v5/_examples"
+	"github.com/go-git/go-git/v5/plumbing/object"
+)
+
+func prepareRepo(w *git.Worktree, directory string) {
+	// We need a known state of files inside the worktree for testing revert a modify and delete
+	Info("echo \"hello world! Modify\" > for-modify")
+	err := ioutil.WriteFile(filepath.Join(directory, "for-modify"), []byte("hello world! Modify"), 0644)
+	CheckIfError(err)
+	Info("git add for-modify")
+	_, err = w.Add("for-modify")
+	CheckIfError(err)
+
+	Info("echo \"hello world! Delete\" > for-delete")
+	err = ioutil.WriteFile(filepath.Join(directory, "for-delete"), []byte("hello world! Delete"), 0644)
+	CheckIfError(err)
+	Info("git add for-delete")
+	_, err = w.Add("for-delete")
+	CheckIfError(err)
+
+	Info("git commit -m \"example go-git commit\"")
+	_, err = w.Commit("example go-git commit", &git.CommitOptions{
+		Author: &object.Signature{
+			Name:  "John Doe",
+			Email: "john@doe.org",
+			When:  time.Now(),
+		},
+	})
+	CheckIfError(err)
+}
+
+// An example of how to restore AKA unstage files
+func main() {
+	CheckArgs("<directory>")
+	directory := os.Args[1]
+
+	// Opens an already existing repository.
+	r, err := git.PlainOpen(directory)
+	CheckIfError(err)
+
+	w, err := r.Worktree()
+	CheckIfError(err)
+
+	prepareRepo(w, directory)
+
+	// Perform the operation and stage them
+	Info("echo \"hello world! Modify 2\" > for-modify")
+	err = ioutil.WriteFile(filepath.Join(directory, "for-modify"), []byte("hello world! Modify 2"), 0644)
+	CheckIfError(err)
+	Info("git add for-modify")
+	_, err = w.Add("for-modify")
+	CheckIfError(err)
+
+	Info("echo \"hello world! Add\" > for-add")
+	err = ioutil.WriteFile(filepath.Join(directory, "for-add"), []byte("hello world! Add"), 0644)
+	CheckIfError(err)
+	Info("git add for-add")
+	_, err = w.Add("for-add")
+	CheckIfError(err)
+
+	Info("rm for-delete")
+	err = os.Remove(filepath.Join(directory, "for-delete"))
+	CheckIfError(err)
+	Info("git add for-delete")
+	_, err = w.Add("for-delete")
+	CheckIfError(err)
+
+	// We can verify the current status of the worktree using the method Status.
+	Info("git status --porcelain")
+	status, err := w.Status()
+	CheckIfError(err)
+	fmt.Println(status)
+
+	// Unstage a single file and see the status
+	Info("git restore --staged for-modify")
+	err = w.Restore(&git.RestoreOptions{Staged: true, Files: []string{"for-modify"}})
+	CheckIfError(err)
+
+	Info("git status --porcelain")
+	status, err = w.Status()
+	CheckIfError(err)
+	fmt.Println(status)
+
+	// Unstage the other 2 files and see the status
+	Info("git restore --staged for-add for-delete")
+	err = w.Restore(&git.RestoreOptions{Staged: true, Files: []string{"for-add", "for-delete"}})
+	CheckIfError(err)
+
+	Info("git status --porcelain")
+	status, err = w.Status()
+	CheckIfError(err)
+	fmt.Println(status)
+}
diff --git a/options.go b/options.go
index d7776da..3cd0f95 100644
--- a/options.go
+++ b/options.go
@@ -416,6 +416,9 @@ type ResetOptions struct {
 	// the index (resetting it to the tree of Commit) and the working tree
 	// depending on Mode. If empty MixedReset is used.
 	Mode ResetMode
+	// Files, if not empty will constrain the reseting the index to only files
+	// specified in this list.
+	Files []string
 }
 
 // Validate validates the fields and sets the default values.
@@ -790,3 +793,26 @@ type PlainInitOptions struct {
 
 // Validate validates the fields and sets the default values.
 func (o *PlainInitOptions) Validate() error { return nil }
+
+var (
+	ErrNoRestorePaths = errors.New("you must specify path(s) to restore")
+)
+
+// RestoreOptions describes how a restore should be performed.
+type RestoreOptions struct {
+	// Marks to restore the content in the index
+	Staged bool
+	// Marks to restore the content of the working tree
+	Worktree bool
+	// List of file paths that will be restored
+	Files []string
+}
+
+// Validate validates the fields and sets the default values.
+func (o *RestoreOptions) Validate() error {
+	if len(o.Files) == 0 {
+		return ErrNoRestorePaths
+	}
+
+	return nil
+}
diff --git a/worktree.go b/worktree.go
index ab11d42..2550da8 100644
--- a/worktree.go
+++ b/worktree.go
@@ -25,11 +25,12 @@ import (
 )
 
 var (
-	ErrWorktreeNotClean     = errors.New("worktree is not clean")
-	ErrSubmoduleNotFound    = errors.New("submodule not found")
-	ErrUnstagedChanges      = errors.New("worktree contains unstaged changes")
-	ErrGitModulesSymlink    = errors.New(gitmodulesFile + " is a symlink")
-	ErrNonFastForwardUpdate = errors.New("non-fast-forward update")
+	ErrWorktreeNotClean                = errors.New("worktree is not clean")
+	ErrSubmoduleNotFound               = errors.New("submodule not found")
+	ErrUnstagedChanges                 = errors.New("worktree contains unstaged changes")
+	ErrGitModulesSymlink               = errors.New(gitmodulesFile + " is a symlink")
+	ErrNonFastForwardUpdate            = errors.New("non-fast-forward update")
+	ErrRestoreWorktreeOnlyNotSupported = errors.New("worktree only is not supported")
 )
 
 // Worktree represents a git worktree.
@@ -307,13 +308,13 @@ func (w *Worktree) ResetSparsely(opts *ResetOptions, dirs []string) error {
 	}
 
 	if opts.Mode == MixedReset || opts.Mode == MergeReset || opts.Mode == HardReset {
-		if err := w.resetIndex(t, dirs); err != nil {
+		if err := w.resetIndex(t, dirs, opts.Files); err != nil {
 			return err
 		}
 	}
 
 	if opts.Mode == MergeReset || opts.Mode == HardReset {
-		if err := w.resetWorktree(t); err != nil {
+		if err := w.resetWorktree(t, opts.Files); err != nil {
 			return err
 		}
 	}
@@ -321,12 +322,47 @@ func (w *Worktree) ResetSparsely(opts *ResetOptions, dirs []string) error {
 	return nil
 }
 
+// Restore restores specified files in the working tree or stage with contents from
+// a restore source. If a path is tracked but does not exist in the restore,
+// source, it will be removed to match the source.
+//
+// If Staged and Worktree are true, then the restore source will be the index.
+// If only Staged is true, then the restore source will be HEAD.
+// If only Worktree is true or neither Staged nor Worktree are true, will
+// result in ErrRestoreWorktreeOnlyNotSupported because restoring the working
+// tree while leaving the stage untouched is not currently supported.
+//
+// Restore with no files specified will return ErrNoRestorePaths.
+func (w *Worktree) Restore(o *RestoreOptions) error {
+	if err := o.Validate(); err != nil {
+		return err
+	}
+
+	if o.Staged {
+		opts := &ResetOptions{
+			Files: o.Files,
+		}
+
+		if o.Worktree {
+			// If we are doing both Worktree and Staging then it is a hard reset
+			opts.Mode = HardReset
+		} else {
+			// If we are doing just staging then it is a mixed reset
+			opts.Mode = MixedReset
+		}
+
+		return w.Reset(opts)
+	}
+
+	return ErrRestoreWorktreeOnlyNotSupported
+}
+
 // Reset the worktree to a specified state.
 func (w *Worktree) Reset(opts *ResetOptions) error {
 	return w.ResetSparsely(opts, nil)
 }
 
-func (w *Worktree) resetIndex(t *object.Tree, dirs []string) error {
+func (w *Worktree) resetIndex(t *object.Tree, dirs []string, files []string) error {
 	idx, err := w.r.Storer.Index()
 	if len(dirs) > 0 {
 		idx.SkipUnless(dirs)
@@ -362,6 +398,13 @@ func (w *Worktree) resetIndex(t *object.Tree, dirs []string) error {
 			name = ch.From.String()
 		}
 
+		if len(files) > 0 {
+			contains := inFiles(files, name)
+			if !contains {
+				continue
+			}
+		}
+
 		b.Remove(name)
 		if e == nil {
 			continue
@@ -379,7 +422,17 @@ func (w *Worktree) resetIndex(t *object.Tree, dirs []string) error {
 	return w.r.Storer.SetIndex(idx)
 }
 
-func (w *Worktree) resetWorktree(t *object.Tree) error {
+func inFiles(files []string, v string) bool {
+	for _, s := range files {
+		if s == v {
+			return true
+		}
+	}
+
+	return false
+}
+
+func (w *Worktree) resetWorktree(t *object.Tree, files []string) error {
 	changes, err := w.diffStagingWithWorktree(true, false)
 	if err != nil {
 		return err
@@ -395,6 +448,25 @@ func (w *Worktree) resetWorktree(t *object.Tree) error {
 		if err := w.validChange(ch); err != nil {
 			return err
 		}
+
+		if len(files) > 0 {
+			file := ""
+			if ch.From != nil {
+				file = ch.From.Name()
+			} else if ch.To != nil {
+				file = ch.To.Name()
+			}
+
+			if file == "" {
+				continue
+			}
+
+			contains := inFiles(files, file)
+			if !contains {
+				continue
+			}
+		}
+
 		if err := w.checkoutChange(ch, t, b); err != nil {
 			return err
 		}
diff --git a/worktree_test.go b/worktree_test.go
index 636ccbe..e1a6c0a 100644
--- a/worktree_test.go
+++ b/worktree_test.go
@@ -25,6 +25,7 @@ import (
 	"github.com/go-git/go-git/v5/storage/memory"
 	"github.com/stretchr/testify/assert"
 
+	"github.com/go-git/go-billy/v5"
 	"github.com/go-git/go-billy/v5/memfs"
 	"github.com/go-git/go-billy/v5/osfs"
 	"github.com/go-git/go-billy/v5/util"
@@ -3053,3 +3054,173 @@ func TestWindowsValidPath(t *testing.T) {
 		})
 	}
 }
+
+var statusCodeNames = map[StatusCode]string{
+	Unmodified:         "Unmodified",
+	Untracked:          "Untracked",
+	Modified:           "Modified",
+	Added:              "Added",
+	Deleted:            "Deleted",
+	Renamed:            "Renamed",
+	Copied:             "Copied",
+	UpdatedButUnmerged: "UpdatedButUnmerged",
+}
+
+func setupForRestore(c *C, s *WorktreeSuite) (fs billy.Filesystem, w *Worktree, names []string) {
+	fs = memfs.New()
+	w = &Worktree{
+		r:          s.Repository,
+		Filesystem: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{})
+	c.Assert(err, IsNil)
+
+	names = []string{"foo", "CHANGELOG", "LICENSE", "binary.jpg"}
+	verifyStatus(c, "Checkout", w, names, []FileStatus{
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Untracked, Staging: Untracked},
+	})
+
+	// Touch of bunch of files including create a new file and delete an exsiting file
+	for _, name := range names {
+		err = util.WriteFile(fs, name, []byte("Foo Bar"), 0755)
+		c.Assert(err, IsNil)
+	}
+	err = util.RemoveAll(fs, names[3])
+	c.Assert(err, IsNil)
+
+	// Confirm the status after doing the edits without staging anything
+	verifyStatus(c, "Edits", w, names, []FileStatus{
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Modified, Staging: Unmodified},
+		{Worktree: Modified, Staging: Unmodified},
+		{Worktree: Deleted, Staging: Unmodified},
+	})
+
+	// Stage all files and verify the updated status
+	for _, name := range names {
+		_, err = w.Add(name)
+		c.Assert(err, IsNil)
+	}
+	verifyStatus(c, "Staged", w, names, []FileStatus{
+		{Worktree: Unmodified, Staging: Added},
+		{Worktree: Unmodified, Staging: Modified},
+		{Worktree: Unmodified, Staging: Modified},
+		{Worktree: Unmodified, Staging: Deleted},
+	})
+
+	// Add secondary changes to a file to make sure we only restore the staged file
+	err = util.WriteFile(fs, names[1], []byte("Foo Bar:11"), 0755)
+	c.Assert(err, IsNil)
+	err = util.WriteFile(fs, names[2], []byte("Foo Bar:22"), 0755)
+	c.Assert(err, IsNil)
+
+	verifyStatus(c, "Secondary Edits", w, names, []FileStatus{
+		{Worktree: Unmodified, Staging: Added},
+		{Worktree: Modified, Staging: Modified},
+		{Worktree: Modified, Staging: Modified},
+		{Worktree: Unmodified, Staging: Deleted},
+	})
+
+	return
+}
+
+func verifyStatus(c *C, marker string, w *Worktree, files []string, statuses []FileStatus) {
+	c.Assert(len(files), Equals, len(statuses))
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+
+	for i, file := range files {
+		current := status.File(file)
+		expected := statuses[i]
+		c.Assert(current.Worktree, Equals, expected.Worktree, Commentf("%s - [%d] : %s Worktree %s != %s", marker, i, file, statusCodeNames[current.Worktree], statusCodeNames[expected.Worktree]))
+		c.Assert(current.Staging, Equals, expected.Staging, Commentf("%s - [%d] : %s Staging %s != %s", marker, i, file, statusCodeNames[current.Staging], statusCodeNames[expected.Staging]))
+	}
+}
+
+func (s *WorktreeSuite) TestRestoreStaged(c *C) {
+	fs, w, names := setupForRestore(c, s)
+
+	// Attempt without files should throw an error like the git restore --staged
+	opts := RestoreOptions{Staged: true}
+	err := w.Restore(&opts)
+	c.Assert(err, Equals, ErrNoRestorePaths)
+
+	// Restore Staged files in 2 groups and confirm status
+	opts.Files = []string{names[0], names[1]}
+	err = w.Restore(&opts)
+	c.Assert(err, IsNil)
+	verifyStatus(c, "Restored First", w, names, []FileStatus{
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Modified, Staging: Unmodified},
+		{Worktree: Modified, Staging: Modified},
+		{Worktree: Unmodified, Staging: Deleted},
+	})
+
+	// Make sure the restore didn't overwrite our secondary changes
+	contents, err := util.ReadFile(fs, names[1])
+	c.Assert(err, IsNil)
+	c.Assert(string(contents), Equals, "Foo Bar:11")
+
+	opts.Files = []string{names[2], names[3]}
+	err = w.Restore(&opts)
+	c.Assert(err, IsNil)
+	verifyStatus(c, "Restored Second", w, names, []FileStatus{
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Modified, Staging: Unmodified},
+		{Worktree: Modified, Staging: Unmodified},
+		{Worktree: Deleted, Staging: Unmodified},
+	})
+
+	// Make sure the restore didn't overwrite our secondary changes
+	contents, err = util.ReadFile(fs, names[2])
+	c.Assert(err, IsNil)
+	c.Assert(string(contents), Equals, "Foo Bar:22")
+}
+
+func (s *WorktreeSuite) TestRestoreWorktree(c *C) {
+	_, w, names := setupForRestore(c, s)
+
+	// Attempt without files should throw an error like the git restore
+	opts := RestoreOptions{}
+	err := w.Restore(&opts)
+	c.Assert(err, Equals, ErrNoRestorePaths)
+
+	opts.Files = []string{names[0], names[1]}
+	err = w.Restore(&opts)
+	c.Assert(err, Equals, ErrRestoreWorktreeOnlyNotSupported)
+}
+
+func (s *WorktreeSuite) TestRestoreBoth(c *C) {
+	_, w, names := setupForRestore(c, s)
+
+	// Attempt without files should throw an error like the git restore --staged --worktree
+	opts := RestoreOptions{Staged: true, Worktree: true}
+	err := w.Restore(&opts)
+	c.Assert(err, Equals, ErrNoRestorePaths)
+
+	// Restore Staged files in 2 groups and confirm status
+	opts.Files = []string{names[0], names[1]}
+	err = w.Restore(&opts)
+	c.Assert(err, IsNil)
+	verifyStatus(c, "Restored First", w, names, []FileStatus{
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Modified, Staging: Modified},
+		{Worktree: Unmodified, Staging: Deleted},
+	})
+
+	opts.Files = []string{names[2], names[3]}
+	err = w.Restore(&opts)
+	c.Assert(err, IsNil)
+	verifyStatus(c, "Restored Second", w, names, []FileStatus{
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Untracked, Staging: Untracked},
+		{Worktree: Untracked, Staging: Untracked},
+	})
+}