git: worktree, Convert from RestoreStaged(files ...string) to Restore…

…(o *RestoreOptions) and add
go-git · ben-tbotlabs · Jun 11, 2021 · Jun 11, 2021 · Jun 11, 2021 · Jul 5, 2021
commit 58302ee411a64b289ecce39eefa01c1824a55050
diff --git a/_examples/restore/main.go b/_examples/restore/main.go
@@ -83,7 +83,7 @@ func main() {
 
 	// Unstage a single file and see the status
 	Info("git restore --staged for-modify")
-	err = w.RestoreStaged("for-modify")
+	err = w.Restore(&git.RestoreOptions{Staged: true, Files: []string{"for-modify"}})
 	CheckIfError(err)
 
 	Info("git status --porcelain")
@@ -93,7 +93,7 @@ func main() {
 
 	// Unstage the other 2 files and see the status
 	Info("git restore --staged for-add for-delete")
-	err = w.RestoreStaged("for-add", "for-delete")
+	err = w.Restore(&git.RestoreOptions{Staged: true, Files: []string{"for-add", "for-delete"}})
 	CheckIfError(err)
 
 	Info("git status --porcelain")

diff --git a/options.go b/options.go
@@ -507,6 +507,7 @@ var (
 	ErrMissingName    = errors.New("name field is required")
 	ErrMissingTagger  = errors.New("tagger field is required")
 	ErrMissingMessage = errors.New("message field is required")
+	ErrNoRestorePaths = errors.New("you must specify path(s) to restore")
 )
 
 // CreateTagOptions describes how a tag object should be created.
@@ -635,3 +636,22 @@ type PlainOpenOptions struct {
 
 // Validate validates the fields and sets the default values.
 func (o *PlainOpenOptions) Validate() error { return nil }
+
+// 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
@@ -26,11 +26,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")
+	ErrRestoreWorktreeeOnlyNotSupported = errors.New("--worktree only is not supported")
 )
 
 // Worktree represents a git worktree.
@@ -263,13 +264,28 @@ func (w *Worktree) setHEADToBranch(branch plumbing.ReferenceName, commit plumbin
 }
 
 // unstages a set of files -- equivalent to "git restore --staged <file>..."
-func (w *Worktree) RestoreStaged(files ...string) error {
-	opts := &ResetOptions{
-		Mode:  MixedReset,
-		Files: files,
+func (w *Worktree) Restore(o *RestoreOptions) error {
+	if err := o.Validate(); err != nil {
+		return err
 	}
 
-	return w.Reset(opts)
+	if o.Worktree && o.Staged {
+		// If we are doing both Worktree and Staging then it is a hard reset
+		opts := &ResetOptions{
+			Mode:  HardReset,
+			Files: o.Files,
+		}
+		return w.Reset(opts)
+	} else if o.Staged {
+		// If we are doing just staging then it is a mixed reset
+		opts := &ResetOptions{
+			Mode:  MixedReset,
+			Files: o.Files,
+		}
+		return w.Reset(opts)
+	} else {
+		return ErrRestoreWorktreeeOnlyNotSupported
+	}
 }
 
 // Reset the worktree to a specified state.
@@ -309,7 +325,7 @@ func (w *Worktree) Reset(opts *ResetOptions) error {
 	}
 
 	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
 		}
 	}
@@ -383,7 +399,7 @@ func inFiles(files []string, v string) bool {
 	return false
 }
 
-func (w *Worktree) resetWorktree(t *object.Tree) error {
+func (w *Worktree) resetWorktree(t *object.Tree, files []string) error {
 	changes, err := w.diffStagingWithWorktree(true)
 	if err != nil {
 		return err
@@ -396,6 +412,20 @@ func (w *Worktree) resetWorktree(t *object.Tree) error {
 	b := newIndexBuilder(idx)
 
 	for _, ch := range changes {
+		if len(files) > 0 {
+			file := ""
+			if ch.From != nil {
+				file = ch.From.Name()
+			} else {
+				file = ch.To.Name()
+			}
+
+			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
@@ -23,6 +23,7 @@ import (
 	"github.com/go-git/go-git/v5/plumbing/object"
 	"github.com/go-git/go-git/v5/storage/memory"
 
+	"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"
@@ -2255,32 +2256,25 @@ var statusCodeNames = map[StatusCode]string{
 	UpdatedButUnmerged: "UpdatedButUnmerged",
 }
 
-func verifyStatus(c *C, 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("[%d] : %s Worktree %s != %s", i, file, statusCodeNames[current.Worktree], statusCodeNames[expected.Worktree]))
-		c.Assert(current.Staging, Equals, expected.Staging, Commentf("[%d] : %s Staging %s != %s", i, file, statusCodeNames[current.Staging], statusCodeNames[expected.Staging]))
-	}
-}
-
-func (s *WorktreeSuite) TestRestoreStaged(c *C) {
-	fs := memfs.New()
-	w := &Worktree{
+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
-	names := []string{"foo", "CHANGELOG", "LICENSE", "binary.jpg"}
 	for i, name := range names {
 		contents := fmt.Sprintf("Foo Bar:%d", i)
 		err = util.WriteFile(fs, name, []byte(contents), 0755)
@@ -2290,39 +2284,134 @@ func (s *WorktreeSuite) TestRestoreStaged(c *C) {
 	c.Assert(err, IsNil)
 
 	//Confirm the status after doing the edits without staging anything
-	verifyStatus(c, w, names, []FileStatus{
+	verifyStatus(c, "Edits", w, names, []FileStatus{
 		{Worktree: Untracked, Staging: Untracked},
 		{Worktree: Modified, Staging: Unmodified},
 		{Worktree: Modified, Staging: Unmodified},
 		{Worktree: Deleted, Staging: Unmodified},
 	})
 
-	// Stage all 3 files and verify the updated status
+	// Stage all files and verify the updated status
 	for _, name := range names {
 		_, err = w.Add(name)
 		c.Assert(err, IsNil)
 	}
-	verifyStatus(c, w, names, []FileStatus{
+	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
-	w.RestoreStaged(names[0], names[1])
-	verifyStatus(c, w, names, []FileStatus{
+	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: Unmodified, Staging: Modified},
+		{Worktree: Modified, Staging: Modified},
 		{Worktree: Unmodified, Staging: Deleted},
 	})
 
-	w.RestoreStaged(names[2], names[3])
-	verifyStatus(c, w, names, []FileStatus{
+	//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, ErrRestoreWorktreeeOnlyNotSupported)
+}
+
+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},
+	})
 }