git: worktree, add RestoreStaged which works like the "git restore --staged <file>..." command

Ben Talbot · Ben Talbot · commit 0f3639790292 · 2024-04-19T22:13:15.000-07:00
Small formatting and style fixes before rebasing against master

Setup args for restore in TestExamples

Fix typo in error message and remove dependency on fmt in worktree_test
diff --git 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
@@ -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
@@ -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
@@ -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,26 +308,61 @@ 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
 		}
 	}
 
 	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
