restore --staged

cmd/restore.go

@@ -8,20 +8,23 @@ import (
 	"os"
 )
 
+var restoreStaged bool
+
 var restoreCmd = &cobra.Command{
 	Use:  "restore",
 	Args: cobra.ExactArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
 		if err := configure(); err != nil {
 			log.Fatalln(err)
 		}
-		if err := mygit.Restore(args[0]); err != nil {
+		if err := mygit.Restore(args[0], restoreStaged); err != nil {
 			fmt.Println(err)
 			os.Exit(1)
 		}
 	},
 }
 
 func init() {
+	restoreCmd.Flags().BoolVar(&restoreStaged, "staged", true, "--staged")
 	rootCmd.AddCommand(restoreCmd)
 }

cmd/status.go

@@ -1,19 +1,23 @@
 package cmd
 
 import (
+	"fmt"
 	"github.com/richardjennings/mygit/internal/mygit"
 	"github.com/spf13/cobra"
-	"log"
 	"os"
 )
 
 var statusCmd = &cobra.Command{
 	Use: "status",
-	RunE: func(cmd *cobra.Command, args []string) error {
+	Run: func(cmd *cobra.Command, args []string) {
 		if err := configure(); err != nil {
-			log.Fatalln(err)
+			fmt.Println(err)
+			os.Exit(1)
+		}
+		if err := mygit.Status(os.Stdout); err != nil {
+			fmt.Println(err)
+			os.Exit(1)
 		}
-		return mygit.Status(os.Stdout)
 	},
 }
 

internal/mygit/index/index.go

@@ -2,6 +2,7 @@ package index
 
 import (
 	"errors"
+	"fmt"
 	"github.com/richardjennings/mygit/internal/mygit/config"
 	"github.com/richardjennings/mygit/internal/mygit/gfs"
 	"os"
@@ -64,6 +65,19 @@ func (idx *Index) File(path string) *gfs.File {
 	return nil
 }
 
+// Rm removes a gfs.File from the Index
+// A call to idx.Write is required to persist the change.
+func (idx *Index) Rm(path string) error {
+	for i, v := range idx.items {
+		if string(v.Name) == path {
+			idx.items = append(idx.items[:i], idx.items[i+1:]...)
+			idx.header.NumEntries--
+			return nil
+		}
+	}
+	return fmt.Errorf("error: pathspec '%s' did not match any file(s) known to git", path)
+}
+
 // Add adds a fs.File to the Index Struct. A call to idx.Write is required
 // to flush the changes to the filesystem.
 func (idx *Index) Add(f *gfs.File) error {

internal/mygit/mygit.go

@@ -401,7 +401,7 @@ func SwitchBranch(name string) error {
 
 }
 
-func Restore(path string) error {
+func Restore(path string, staged bool) error {
 	idx, err := index.ReadIndex()
 	if err != nil {
 		return err
@@ -414,6 +414,14 @@ func Restore(path string) error {
 	if err != nil {
 		return err
 	}
+	if staged {
+		// remove file from index
+		if err := idx.Rm(path); err != nil {
+			return err
+		}
+		return idx.Write()
+	}
+
 	fileStatus, ok := currentStatus.Contains(path)
 	// if the path not found or is untracked working directory fileStatus then error
 	if !ok || fileStatus.WdStatus == gfs.WDUntracked {
@@ -454,5 +462,5 @@ func Restore(path string) error {
 	if err := fh.Close(); err != nil {
 		return err
 	}
-	return os.Chtimes(path, file.Finfo.ModTime(), file.Finfo.ModTime())
+	return os.Chtimes(filepath.Join(config.Path(), path), file.Finfo.ModTime(), file.Finfo.ModTime())
 }

internal/mygit/mygit_test.go

@@ -156,6 +156,23 @@ func Test_End_To_End(t *testing.T) {
 	// switch back to main, should get file back
 	testSwitchBranch(t, "main")
 	testStatus(t, "")
+
+	// test restore staged
+	writeFile(t, dir, "o", []byte("o"))
+	testAdd(t, "o", 3)
+	testStatus(t, "A  o\n")
+	testRestore(t, "o", true)
+	testStatus(t, "?? o\n")
+
+	// test restore
+	testAdd(t, "o", 3)
+	testCommit(t, []byte("oo"))
+	testStatus(t, "")
+	writeFile(t, dir, "o", []byte("ok"))
+	testStatus(t, " M o\n")
+	testRestore(t, "o", false)
+	testStatus(t, "")
+
 }
 
 func testDir(t *testing.T) string {
@@ -212,6 +229,12 @@ func testStatus(t *testing.T, expected string) {
 	assert.Equal(t, expected, buf.String())
 }
 
+func testRestore(t *testing.T, path string, staged bool) {
+	if err := Restore(path, staged); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func testCommit(t *testing.T, message []byte) []byte {
 	sha, err := Commit(message)
 	if err != nil {