Merge pull request #361 from mcuadros/worktree-add

worktree: add method

3 files changed, 205 insertions, 4 deletions

diff --git a/worktree_linux.go b/worktree_linux.go
index a33cd2f..7209d7d 100644
--- a/worktree_linux.go
+++ b/worktree_linux.go
@@ -12,6 +12,7 @@ import (
 func init() {
 	fillSystemInfo = func(e *index.Entry, sys interface{}) {
 		if os, ok := sys.(*syscall.Stat_t); ok {
+
 			e.CreatedAt = time.Unix(int64(os.Ctim.Sec), int64(os.Ctim.Nsec))
 			e.Dev = uint32(os.Dev)
 			e.Inode = uint32(os.Ino)
diff --git a/worktree_status.go b/worktree_status.go
index 373f161..6becada 100644
--- a/worktree_status.go
+++ b/worktree_status.go
@@ -2,16 +2,20 @@ package git
 
 import (
 	"bytes"
+	"io"
 
 	"gopkg.in/src-d/go-git.v4/plumbing"
+	"gopkg.in/src-d/go-git.v4/plumbing/filemode"
+	"gopkg.in/src-d/go-git.v4/plumbing/format/index"
 	"gopkg.in/src-d/go-git.v4/plumbing/object"
+	"gopkg.in/src-d/go-git.v4/utils/ioutil"
 	"gopkg.in/src-d/go-git.v4/utils/merkletrie"
 	"gopkg.in/src-d/go-git.v4/utils/merkletrie/filesystem"
-	"gopkg.in/src-d/go-git.v4/utils/merkletrie/index"
+	mindex "gopkg.in/src-d/go-git.v4/utils/merkletrie/index"
 	"gopkg.in/src-d/go-git.v4/utils/merkletrie/noder"
 )
 
-// Status returns the working tree status
+// Status returns the working tree status.
 func (w *Worktree) Status() (Status, error) {
 	ref, err := w.r.Head()
 	if err == plumbing.ErrReferenceNotFound {
@@ -80,7 +84,7 @@ func (w *Worktree) diffStagingWithWorktree() (merkletrie.Changes, error) {
 		return nil, err
 	}
 
-	from := index.NewRootNode(idx)
+	from := mindex.NewRootNode(idx)
 	submodules, err := w.getSubmodulesStatus()
 	if err != nil {
 		return nil, err
@@ -131,7 +135,7 @@ func (w *Worktree) diffCommitWithStaging(commit plumbing.Hash, reverse bool) (me
 		return nil, err
 	}
 
-	to := index.NewRootNode(idx)
+	to := mindex.NewRootNode(idx)
 	from := object.NewTreeRootNode(t)
 
 	if reverse {
@@ -159,3 +163,105 @@ func diffTreeIsEquals(a, b noder.Hasher) bool {
 
 	return bytes.Equal(hashA, hashB)
 }
+
+// Add adds the file contents of a file in the worktree to the index. if the
+// file is already stagged in the index no error is returned.
+func (w *Worktree) Add(path string) (plumbing.Hash, error) {
+	s, err := w.Status()
+	if err != nil {
+		return plumbing.ZeroHash, err
+	}
+
+	h, err := w.calculateBlobHash(path)
+	if err != nil {
+		return h, err
+	}
+
+	if s.File(path).Worktree == Unmodified {
+		return h, nil
+	}
+
+	if err := w.addOrUpdateFileToIndex(path, h); err != nil {
+		return h, err
+	}
+
+	return h, err
+}
+
+func (w *Worktree) calculateBlobHash(filename string) (hash plumbing.Hash, err error) {
+	fi, err := w.fs.Stat(filename)
+	if err != nil {
+		return plumbing.ZeroHash, err
+	}
+
+	f, err := w.fs.Open(filename)
+	if err != nil {
+		return plumbing.ZeroHash, err
+	}
+
+	defer ioutil.CheckClose(f, &err)
+
+	h := plumbing.NewHasher(plumbing.BlobObject, fi.Size())
+	if _, err := io.Copy(h, f); err != nil {
+		return plumbing.ZeroHash, err
+	}
+
+	hash = h.Sum()
+	return
+}
+
+func (w *Worktree) addOrUpdateFileToIndex(filename string, h plumbing.Hash) error {
+	idx, err := w.r.Storer.Index()
+	if err != nil {
+		return err
+	}
+
+	_, err = idx.Entry(filename)
+	if err == index.ErrEntryNotFound {
+		err = w.doAddFileToIndex(idx, filename)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	err = w.doUpdateFileToIndex(idx, filename, h)
+	if err != nil {
+		return err
+	}
+
+	return w.r.Storer.SetIndex(idx)
+}
+
+func (w *Worktree) doAddFileToIndex(idx *index.Index, filename string) error {
+	idx.Entries = append(idx.Entries, index.Entry{
+		Name: filename,
+	})
+
+	return nil
+}
+
+func (w *Worktree) doUpdateFileToIndex(idx *index.Index, filename string, h plumbing.Hash) error {
+	info, err := w.fs.Stat(filename)
+	if err != nil {
+		return err
+	}
+
+	for i, e := range idx.Entries {
+		if e.Name != filename {
+			continue
+		}
+
+		e.Hash = h
+		e.ModifiedAt = info.ModTime()
+		e.Mode, err = filemode.NewFromOSFileMode(info.Mode())
+		if err != nil {
+			return err
+		}
+
+		fillSystemInfo(&e, info.Sys())
+		idx.Entries[i] = e
+	}
+
+	return nil
+}
diff --git a/worktree_test.go b/worktree_test.go
index 3393469..f06a1f9 100644
--- a/worktree_test.go
+++ b/worktree_test.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/src-d/go-git-fixtures"
 	. "gopkg.in/check.v1"
+	"gopkg.in/src-d/go-billy.v2"
 	"gopkg.in/src-d/go-billy.v2/memfs"
 	"gopkg.in/src-d/go-billy.v2/osfs"
 )
@@ -478,3 +479,96 @@ func (s *WorktreeSuite) TestSubmodules(c *C) {
 
 	c.Assert(l, HasLen, 2)
 }
+
+func (s *WorktreeSuite) TestAddUntracked(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:  s.Repository,
+		fs: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	idx, err := w.r.Storer.Index()
+	c.Assert(err, IsNil)
+	c.Assert(idx.Entries, HasLen, 9)
+
+	err = billy.WriteFile(w.fs, "foo", []byte("FOO"), 0755)
+	c.Assert(err, IsNil)
+
+	hash, err := w.Add("foo")
+	c.Assert(hash.String(), Equals, "d96c7efbfec2814ae0301ad054dc8d9fc416c9b5")
+	c.Assert(err, IsNil)
+
+	idx, err = w.r.Storer.Index()
+	c.Assert(err, IsNil)
+	c.Assert(idx.Entries, HasLen, 10)
+
+	e, err := idx.Entry("foo")
+	c.Assert(err, IsNil)
+	c.Assert(e.Hash, Equals, hash)
+	c.Assert(e.Mode, Equals, filemode.Executable)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 1)
+
+	file := status.File("foo")
+	c.Assert(file.Staging, Equals, Added)
+	c.Assert(file.Worktree, Equals, Unmodified)
+}
+
+func (s *WorktreeSuite) TestAddModified(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:  s.Repository,
+		fs: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	idx, err := w.r.Storer.Index()
+	c.Assert(err, IsNil)
+	c.Assert(idx.Entries, HasLen, 9)
+
+	err = billy.WriteFile(w.fs, "LICENSE", []byte("FOO"), 0644)
+	c.Assert(err, IsNil)
+
+	hash, err := w.Add("LICENSE")
+	c.Assert(err, IsNil)
+	c.Assert(hash.String(), Equals, "d96c7efbfec2814ae0301ad054dc8d9fc416c9b5")
+
+	idx, err = w.r.Storer.Index()
+	c.Assert(err, IsNil)
+	c.Assert(idx.Entries, HasLen, 9)
+
+	e, err := idx.Entry("LICENSE")
+	c.Assert(err, IsNil)
+	c.Assert(e.Hash, Equals, hash)
+	c.Assert(e.Mode, Equals, filemode.Regular)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 1)
+
+	file := status.File("LICENSE")
+	c.Assert(file.Staging, Equals, Modified)
+	c.Assert(file.Worktree, Equals, Unmodified)
+}
+
+func (s *WorktreeSuite) TestAddUnmodified(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:  s.Repository,
+		fs: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	hash, err := w.Add("LICENSE")
+	c.Assert(hash.String(), Equals, "c192bd6a24ea1ab01d78686e417c8bdc7c3d197f")
+	c.Assert(err, IsNil)
+}